Ë
    tPêj  ã                  óL   — d Z ddlmZ ddlmZ ddlmZmZ ddlZ G d„ d«      Z	y)z)Database access helpers for Focus export.é    )Úannotations)Údatetime)ÚAnyÚFinalNc                  ó>   — e Zd ZdZd„ Zddddœ	 	 	 	 	 	 	 dd„Zdd„Zy)	ÚFocusLiteLLMDatabasez8Retrieves LiteLLM usage data for Focus export workflows.c                ó,   — ddl m} |€t        d«      ‚|S )Nr   )Úprisma_clientz�Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys)Úlitellm.proxy.proxy_serverr
   ÚRuntimeError)Úselfr
   s     úh/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/integrations/focus/database.pyÚ_ensure_prisma_clientz*FocusLiteLLMDatabase._ensure_prisma_client   s&   € Ý<àÐ ÜðYóð ð Ðó    N)ÚlimitÚstart_time_utcÚend_time_utcc             ƒ  óœ  K  — | j                  «       }g }g }d}|r+|j                  d|› d�«       |j                  |«       |dz  }|r+|j                  d|› d�«       |j                  |«       |dz  }d}|rddj                  |«      z   }d}	|�2	 t        |«      }
|
d
k  rt        d«      ‚d|› �}	|j                  |
«       d|› d|	› d�}	  |j                  j                  |g|¢­Ž ƒ d{  –—† }t        j                  |d¬«      S # t        t
        f$ r}t        d	«      |‚d}~ww xY w7 Œ># t        $ r}t        d|› �«      |‚d}~ww xY w­w)z+Return usage data for the requested window.é   zdus.updated_at >= $z::timestamptzzdus.updated_at <= $Ú zWHERE z AND Nzlimit must be an integerr   zlimit must be non-negativez LIMIT $a¬  
        SELECT
            dus.id,
            dus.date,
            dus.user_id,
            dus.api_key,
            dus.model,
            dus.model_group,
            dus.custom_llm_provider,
            dus.prompt_tokens,
            dus.completion_tokens,
            dus.spend,
            dus.api_requests,
            dus.successful_requests,
            dus.failed_requests,
            dus.cache_creation_input_tokens,
            dus.cache_read_input_tokens,
            dus.created_at,
            dus.updated_at,
            vt.team_id,
            vt.key_alias as api_key_alias,
            tt.team_alias,
            ut.user_email as user_email,
            COALESCE(vt.organization_id, tt.organization_id) as organization_id,
            ot.organization_alias as organization_alias
        FROM "LiteLLM_DailyUserSpend" dus
        LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
        LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
        LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
        LEFT JOIN "LiteLLM_OrganizationTable" ot
            ON ot.organization_id = COALESCE(vt.organization_id, tt.organization_id)
        z=
        ORDER BY dus.date DESC, dus.created_at DESC
        z	
        )Úinfer_schema_lengthzError retrieving usage data: )r   ÚappendÚjoinÚintÚ	TypeErrorÚ
ValueErrorÚdbÚ	query_rawÚplÚ	DataFrameÚ	Exceptionr   )r   r   r   r   ÚclientÚwhere_clausesÚquery_paramsÚplaceholder_indexÚwhere_clauseÚlimit_clauseÚlimit_valueÚexcÚqueryÚdb_responses                 r   Úget_usage_dataz#FocusLiteLLMDatabase.get_usage_data   s¦  è ø€ ð ×2Ñ2Ó4ˆà*,ˆØ)+ˆØÐÙØ× Ñ Ð#6Ð7HÐ6IÈÐ!WÔXØ×Ñ Ô/Ø Ñ"ÐÙØ× Ñ Ð#6Ð7HÐ6IÈÐ!WÔXØ×Ñ Ô-Ø Ñ"ÐàˆÙØ# g§l¡l°=Ó&AÑAˆLàˆØÐðFÜ%(¨£Z�ð ˜QŠÜ Ð!=Ó>Ð>Ø%Ð&7Ð%8Ð9ˆLØ×Ñ Ô,ð	ð> 
ˆð 	à	ˆð 	ðC"ˆðH	OØ': v§y¡y×':Ñ':¸5Ð'PÀ<Ò'P×!PˆKÜ—<‘< ÀÔFÐFøô[ œzÐ*ò FÜ Ð!;Ó<À#ÐEûðFúðX "Qùäò 	OÜÐ!>¸s¸eÐDÓEÈ3ÐNûð	Oüs`   ‚BEÂD Â/EÃ D, Ã+D*Ã,D, ÄEÄD'ÄD"Ä"D'Ä'EÄ*D, Ä,	E	Ä5EÅE	Å	Ec              ƒ  óÄ   K  — | j                  «       }d}	 |j                  j                  |«      ƒ d{  –—† }|ddœS 7 Œ	# t        $ r}t	        d|› �«      |‚d}~ww xY w­w)z6Return metadata about the spend table for diagnostics.z»
        SELECT column_name, data_type, is_nullable
        FROM information_schema.columns
        WHERE table_name = 'LiteLLM_DailyUserSpend'
        ORDER BY ordinal_position;
        NÚLiteLLM_DailyUserSpend)ÚcolumnsÚ
table_namezError getting table info: )r   r   r   r!   r   )r   r"   Ú
info_queryÚcolumns_responser)   s        r   Úget_table_infoz#FocusLiteLLMDatabase.get_table_infog   sq   è ø€ à×2Ñ2Ó4ˆðˆ
ð	LØ,2¯I©I×,?Ñ,?À
Ó,K×&KÐØ/Ð?WÑXÐXð 'Lùäò 	LÜÐ!;¸C¸5ÐAÓBÈÐKûð	Lüs5   ‚A –A  ´>µA  ½A ¾A  Á 	AÁ	AÁAÁA )r   z
int | Noner   údatetime | Noner   r4   Úreturnzpl.DataFrame)r5   zdict[str, Any])Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   r,   r3   © r   r   r   r      sS   „ ÙBòð !Ø*.Ø(,ñMOð ðMOð (ð	MOð
 &ðMOð 
óMOô^Lr   r   )
r9   Ú
__future__r   r   Útypingr   r   Úpolarsr   r   r:   r   r   ú<module>r>      s"   ðÙ /å "å ß ã ÷jLò jLr   