o
    ưig                     @   sh  d Z ddlZddlmZmZmZ ddlmZmZm	Z	m
Z
 ddlmZ ddlmZ ddlT ddlmZ dd	lmZ dd
lmZ ddlmZmZ e Z		d6dedee dededee defddZdedefddZdedee e!f ddfddZ"de dee dededef
ddZ#de dee dee fddZ$d e%de%fd!d"Z&ej'd#d$geege(d%eeefde d&e
defd'd(Z)ej'd)d$geege*d%eeefde!d&e
defd*d+Z+ej,d,d$geegee* d%eeefde-d&e
defd-d.Z.ej/d/d$geege*d%eefd0edefd1d2Z0ej/d3d$geegee* d%eefdefd4d5Z1dS )7zm
Endpoints for /project operations

/project/new
/project/update
/project/delete
/project/info
/project/list
    N)ListOptionalUnion)	APIRouterDependsHTTPExceptionRequest)verbose_proxy_logger)uuid)*)user_api_key_auth)_set_object_metadata_field)management_endpoint_wrapper)PrismaClienthandle_exception_on_proxyFuser_api_key_dictteam_idprisma_clientrequire_adminteam_objectreturnc                    sl   | j tjk}|r|S |rdS |r| jsdS |}|du r)|jjjd|idI dH }|r4|jr4| j|jv S dS )ab  
    Check if user has permission to manage a project.

    Returns True if user is proxy admin or team admin (when team_id provided).
    If require_admin=True, only proxy admins are allowed.

    If team_object is provided, it will be used instead of fetching from DB
    (avoids duplicate DB queries when team was already fetched for validation).
    TFNr   where)	user_roleLitellmUserRolesPROXY_ADMINuser_iddblitellm_teamtablefind_uniqueadmins)r   r   r   r   r   Zis_proxy_adminteam r"   k/home/app/Keep/.python/lib/python3.10/site-packages/litellm/proxy/management_endpoints/project_endpoints.py"_check_user_permission_for_project   s    

r$   c                    s>   |j jjd| idI dH }|du rtd|  dddd|S )z2Validate that a team exists. Returns the team row.r   r   NzTeam not found, team_id=	not_found  messagetypecodeparam)r   r   r   ProxyException)r   r   r!   r"   r"   r#   _validate_team_existsF   s   r-   datac              
   C   s  |j dur|j dk rtddd|j  id|jdur,|jdk r,tddd|j id|jdurM|j durM|j|j krMtddd|j d	|j  d
idt|dd}| jpWg }|rt|dkrtjj|vr|D ]}||vrtddd| d| d| j	 idqh|j dur| j dur|j | j krtddd|j  d| j  d| j	 id|j
dur| j
dur|j
| j
krtddd|j
 d| j
 d| j	 id|jdur| jdur|j| jkrtddd|j d| j d| j	 iddS dS dS )a  
    Check that project limits respect its parent Team's limits.

    Mirrors _check_org_team_limits() from team_endpoints.py.

    Validates:
    - Project models are a subset of Team models
    - Project max_budget <= Team max_budget
    - Project tpm_limit <= Team tpm_limit
    - Project rpm_limit <= Team rpm_limit
    - Budget values are non-negative
    - soft_budget < max_budget
    Nr     errorz)max_budget cannot be negative. Received: status_codedetailz*soft_budget cannot be negative. Received: zsoft_budget (z*) must be strictly lower than max_budget ()modelszModel 'z4' not in team's allowed models. Team allowed models=z. Team: zProject max_budget (z) exceeds team's max_budget (z	). Team: zProject tpm_limit (z) exceeds team's tpm_limit (zProject rpm_limit (z) exceeds team's rpm_limit ()Z
max_budgetr   Zsoft_budgetgetattrr5   lenZSpecialModelNamesZall_proxy_modelsvaluer   Z	tpm_limitZ	rpm_limit)r   r.   Zproject_modelsZteam_modelsmr"   r"   r#   _check_team_project_limitsZ   st   



	

	
r:   r   litellm_proxy_admin_namec           	         s   t j  | jdd} fdd| D }t di |}||jdd}|jjji ||p1||p4|ddI dH }|j	S )	z5Create a budget for the project and return budget_id.Texclude_nonec                       i | ]\}}| v r||qS r"   r"   .0kvZbudget_paramsr"   r#   
<dictcomp>       z._create_budget_for_project.<locals>.<dictcomp>)
created_by
updated_byr.   Nr"   )
LiteLLM_BudgetTablemodel_fieldskeysjsonitemsjsonify_objectr   litellm_budgettablecreate	budget_id)	r.   r   r;   r   Z
_json_dataZ_budget_dataZ
budget_rowZ
new_budgetZ_budgetr"   rC   r#   _create_budget_for_project   s   
rR   c                    sF   |du rdS | j dur!|jjj| j jdddI dH }| ` |jS dS )z
    Creates the LiteLLM_ObjectPermissionTable record for the project.
    Returns the object_permission_id if created, otherwise None.
    NTr<   rH   )object_permissionr   litellm_objectpermissiontablerP   
model_dumpobject_permission_id)r.   r   Zcreated_object_permissionr"   r"   r#   _set_project_object_permission   s   

rW   project_datac                 C   s0   t j }t|D ]}|dkr| |d q	| S )z
    Remove budget fields from project data.
    Budget fields belong to LiteLLM_BudgetTable, not LiteLLM_ProjectTable.
    Keep budget_id as it's a foreign key.

    Following the pattern from organization_endpoints.py
    rQ   N)rI   rJ   rK   listpop)rX   budget_fieldsfieldr"   r"   r#   '_remove_budget_fields_from_project_data   s   
r]   z/project/newzproject management)tagsdependenciesZresponse_modelhttp_requestc              
      s  ddl m}m}m} z)t| dddur#|s#tdddtjj id|s1tddd	tjj idt	D ]}t| |ddurLt
| |t| |d
 t| | q3|du r[tddtjjidt| j|dI dH }ttd"i | | d t|| j|td"i | dI dH }|stddd|j id| jdu rtt | _n|jjjd| jidI dH }	|	durtd| j ddddd| jdu rt| |j||dI dH | _t| |dI dH }
t d"i | j!dd|
|jp||jp|d}t"D ]}t| |ddurt
||t| |d
 q|#|j!dd}t$|}t%&dt!j'|dd  |jjj(i |ddid I dH }|W S  t)yP } zt%*d!+t| t,|d}~ww )#aE  
    Create a new project. Projects sit between teams and keys in the hierarchy.

    Only admins or team admins can create projects.

    # Parameters

    - project_alias: *Optional[str]* - The name of the project.
    - description: *Optional[str]* - Description of the project's purpose and use case.
    - team_id: *str* - The team id that this project belongs to. Required.
    - models: *List* - The models the project has access to.
    - budget_id: *Optional[str]* - The id for a budget (tpm/rpm/max budget) for the project.
    ### IF NO BUDGET ID - CREATE ONE WITH THESE PARAMS ###
    - max_budget: *Optional[float]* - Max budget for project
    - tpm_limit: *Optional[int]* - Max tpm limit for project
    - rpm_limit: *Optional[int]* - Max rpm limit for project
    - max_parallel_requests: *Optional[int]* - Max parallel requests for project
    - soft_budget: *Optional[float]* - Get a slack alert when this soft budget is reached. Don't block requests.
    - model_max_budget: *Optional[dict]* - Max budget for a specific model. Example: {"gpt-4": 100.0, "gpt-3.5-turbo": 50.0}
    - model_rpm_limit: *Optional[dict]* - RPM limits per model. Example: {"gpt-4": 1000, "gpt-3.5-turbo": 5000}
    - model_tpm_limit: *Optional[dict]* - TPM limits per model. Example: {"gpt-4": 50000, "gpt-3.5-turbo": 100000}
    - budget_duration: *Optional[str]* - Frequency of reseting project budget
    - metadata: *Optional[dict]* - Metadata for project, store information for project. Example metadata - {"use_case_id": "SNOW-12345", "responsible_ai_id": "RAI-67890"}
    - tags: *Optional[list]* - Tags for the project. Example: ["production", "api"]
    - blocked: *bool* - Flag indicating if the project is blocked or not - will stop all calls from keys with this project_id.
    - object_permission: Optional[LiteLLM_ObjectPermissionBase] - project-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.

    Example 1: Create new project **without** a budget_id, with model-specific limits

    ```bash
    curl --location 'http://0.0.0.0:4000/project/new' \
    --header 'Authorization: Bearer sk-1234' \
    --header 'Content-Type: application/json' \
    --data '{
        "project_alias": "flight-search-assistant",
        "description": "AI-powered flight search and booking assistant",
        "team_id": "team-123",
        "models": ["gpt-4", "gpt-3.5-turbo"],
        "max_budget": 100,
        "model_rpm_limit": {
            "gpt-4": 1000,
            "gpt-3.5-turbo": 5000
        },
        "model_tpm_limit": {
            "gpt-4": 50000,
            "gpt-3.5-turbo": 100000
        },
        "metadata": {
            "use_case_id": "SNOW-12345",
            "responsible_ai_id": "RAI-67890"
        }
    }'
    ```

    Example 2: Create new project **with** a budget_id

    ```bash
    curl --location 'http://0.0.0.0:4000/project/new' \
    --header 'Authorization: Bearer sk-1234' \
    --header 'Content-Type: application/json' \
    --data '{
        "project_alias": "hotel-recommendations",
        "description": "Personalized hotel recommendation engine",
        "team_id": "team-123",
        "models": ["claude-3-sonnet"],
        "budget_id": "428eeaa8-f3ac-4e85-a8fb-7dc8d7aa8689",
        "metadata": {
            "use_case_id": "SNOW-54321"
        }
    }'
    ```
    r   r;   premium_userr   r^   N  r0   -Only premium users can add tags to projects. r1   -Project management is an enterprise feature. Zobject_data
field_namer8     r   r   r   r.   r   r   r   r   z=Only admins or team admins can create projects. Your role is 
project_idr   zProject id = z3 already exists. Please use a different project id.bad_requestr/   r'   )r.   r   r;   r   )r.   r   Tr<   )rV   rF   rG   znew_project_row:    )indentlitellm_budget_table)r.   includezZlitellm.proxy.management_endpoints.project_endpoints.new_project(): Exception occured - {}r"   )-litellm.proxy.proxy_serverr;   rb   r   r6   r   CommonProxyErrorsnot_premium_userr8   1LiteLLM_ManagementEndpoint_MetadataFields_Premiumr   delattrdb_not_connected_errorr-   r   r:   LiteLLM_TeamTablerU   r$   r   rl   strr
   uuid4r   litellm_projecttabler   r,   rQ   rR   r   rW   LiteLLM_ProjectTablerL   )LiteLLM_ManagementEndpoint_MetadataFieldsrN   r]   r	   infodumpsrP   	Exception	exceptionformatr   )r.   r`   r   r;   rb   r   r\   r   has_permissionexisting_projectrV   Zproject_rowZnew_project_rowresponseer"   r"   r#   new_project   s   T	






r   z/project/updatec              
      sJ  ddl m}m}m} z}t| dddur#|s#tdddtjj id|s1tddd	tjj idt	D ]}t| |ddurLt
| |t| |d
 t| | q3|du r[tddtjjid| jdu rhtdddid|jjjd| jidI dH }|du rtd| j dddd| jp|j}d}	|durt||dI dH }	t||j||	rtd&i |	 nddI dH }
|
stdddid|	durttd&i |	 | d | jddhd}||}|jp||d< tj   fdd| D }|r|jr|jj j!d|jii |d|jp	|idI dH  | D ]	}|"|d qd|v rP|"d}|rP|j#r@|jj$j!d |j#i|dI dH  n|jj$j%|d!I dH }|j#|d < t&D ]}||v rn|'d"du rei |d"< |"||d" |< qRt(|}|jjj!d| ji|ddd#d$I dH }|W S  t)y } zt*+d%,t-| t.|d}~ww )'a;  
    Update a project

    Parameters:
    - project_id: *str* - The project id to update. Required.
    - project_alias: *Optional[str]* - Updated name for the project
    - description: *Optional[str]* - Updated description for the project
    - team_id: *Optional[str]* - Updated team_id for the project
    - metadata: *Optional[dict]* - Updated metadata for project
    - models: *Optional[list]* - Updated list of models for the project
    - blocked: *Optional[bool]* - Updated blocked status
    - max_budget: *Optional[float]* - Updated max budget
    - tpm_limit: *Optional[int]* - Updated tpm limit
    - rpm_limit: *Optional[int]* - Updated rpm limit
    - model_rpm_limit: *Optional[dict]* - Updated RPM limits per model
    - model_tpm_limit: *Optional[dict]* - Updated TPM limits per model
    - budget_duration: *Optional[str]* - Updated budget duration
    - tags: *Optional[list]* - Updated list of tags for the project
    - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Updated object permission

    Example:
    ```bash
    curl --location 'http://0.0.0.0:4000/project/update' \
    --header 'Authorization: Bearer sk-1234' \
    --header 'Content-Type: application/json' \
    --data '{
        "project_id": "project-123",
        "description": "Updated flight search system with enhanced capabilities",
        "max_budget": 200,
        "model_rpm_limit": {
            "gpt-4": 2000,
            "gpt-3.5-turbo": 10000
        },
        "metadata": {
            "use_case_id": "SNOW-12345",
            "status": "active"
        }
    }'
    ```
    r   ra   r^   Nrc   r0   rd   r1   re   rf   rh   r/   zproject_id is requiredrl   r   Project not found, project_id=r%   r&   r'   ri   rk   z.Only admins or team admins can update projectsrj   T)r=   excluderG   c                    r>   r"   r"   r?   r[   r"   r#   rD   w  rE   z"update_project.<locals>.<dictcomp>rQ   )r   r.   rS   rV   rH   metadatarp   rS   )r   r.   rq   z]litellm.proxy.management_endpoints.project_endpoints.update_project(): Exception occured - {}r"   )/rr   r;   rb   r   r6   r   rs   rt   r8   ru   r   rv   rw   rl   r   r{   r   r,   r   r-   r$   rx   rU   r:   rL   rN   r   rI   rJ   rK   rM   rQ   rO   updaterZ   rV   rT   rP   r}   getr]   r   r	   r   r   ry   r   )r.   r`   r   r;   rb   r   r\   r   Zteam_id_to_checkZteam_obj_for_checksr   Zupdate_dataZbudget_updatesZobject_permission_dataZcreated_permissionZupdated_projectr   r"   r   r#   update_project  s   4	



	



	

r   z/project/deletec              
      sd  ddl m}m} z|stdddtjj id|du r&tddtjjidt|d|d	d
I dH }|s;tdddidg }| j	D ]T}|j
jjd|idI dH }|du r^td| dddd|j
jjd|idI dH }	t|	dkrtd| dt|	 ddddd|j
jjd|idI dH }
||
 q@|W S  ty } ztdt| t|d}~ww )a  
    Delete projects

    Parameters:
    - project_ids: *List[str]* - List of project ids to delete

    Example:
    ```bash
    curl --location --request DELETE 'http://0.0.0.0:4000/project/delete' \
    --header 'Authorization: Bearer sk-1234' \
    --header 'Content-Type: application/json' \
    --data '{
        "project_ids": ["project-123", "project-456"]
    }'
    ```
    r   )rb   r   rc   r0   re   r1   Nrh   T)r   r   r   r   zOnly admins can delete projectsrl   r   r   r%   r&   project_idsr'   zCannot delete project z. zI key(s) are associated with it. Please delete or reassign the keys first.rm   r/   z]litellm.proxy.management_endpoints.project_endpoints.delete_project(): Exception occured - {})rr   rb   r   r   rs   rt   r8   rw   r$   r   r   r{   r   r,   Zlitellm_verificationtoken	find_manyr7   deleteappendr   r	   r   r   ry   r   )r.   r`   r   rb   r   r   Zdeleted_projectsrl   r   Zassociated_keysZdeleted_projectr   r"   r"   r#   delete_project  s   

	
r   z/project/inforl   c              
      s  ddl m} zg|du rtddtjjid|jjjd| iddd	d
I dH }|du r6t	d|  dddd|j
tjk}d}|jr`|jr`|jjjd|jidI dH }|r`|j|jv p_|j|jv }|sl|sltdddid|W S  ty } ztdt| t|d}~ww )a  
    Get information about a specific project

    Parameters:
    - project_id: *str* - The project id to fetch info for

    Example:
    ```bash
    curl --location 'http://0.0.0.0:4000/project/info?project_id=project-123' \
    --header 'Authorization: Bearer sk-1234'
    ```
    r   r   Nrh   r0   r1   rl   Tr   r   rq   r   r%   r&   r'   Fr   r   rc   z%You don't have access to this projectz[litellm.proxy.management_endpoints.project_endpoints.project_info(): Exception occured - {})rr   r   r   rs   rw   r8   r   r{   r   r,   r   r   r   r   r   r   r    membersr   r	   r   r   ry   r   )rl   r   r   projectZis_adminZis_team_memberr!   r   r"   r"   r#   project_info$  sX   

r   z/project/listc              
      s   ddl m} z[|du rtddtjjid| jtjkr-|j	j
jdddd	I dH }|W S |j	jjd
dd| jiidd| jiigidI dH }dd |D }|j	j
jdd|iiddddI dH }|W S  ty| } ztdt| t|d}~ww )z
    List all projects that the user has access to

    Example:
    ```bash
    curl --location 'http://0.0.0.0:4000/project/list' \
    --header 'Authorization: Bearer sk-1234'
    ```
    r   r   Nrh   r0   r1   Tr   )rq   ORr   hasr    r   c                 S   s   g | ]}|j qS r"   )r   )r@   r!   r"   r"   r#   
<listcomp>  s    z!list_projects.<locals>.<listcomp>r   inr   z\litellm.proxy.management_endpoints.project_endpoints.list_projects(): Exception occured - {})rr   r   r   rs   rw   r8   r   r   r   r   r{   r   r   r   r   r	   r   r   ry   r   )r   r   ZprojectsZ
user_teamsZteam_idsr   r"   r"   r#   list_projectso  sF   
	
r   )FN)2__doc__rL   typingr   r   r   Zfastapir   r   r   r   Zlitellm._loggingr	   Zlitellm._uuidr
   Zlitellm.proxy._typesZ$litellm.proxy.auth.user_api_key_authr   Z/litellm.proxy.management_endpoints.common_utilsr   Z&litellm.proxy.management_helpers.utilsr   Zlitellm.proxy.utilsr   r   ZrouterZUserAPIKeyAuthry   boolrx   r$   r-   ZNewProjectRequestZUpdateProjectRequestr:   rR   rW   dictr]   postZNewProjectResponser   r|   r   r   ZDeleteProjectRequestr   r   r   r   r"   r"   r"   r#   <module>   s    
'


d

 ] SeE