
    j?                        d dl Z d dlmZ ddlmZ ddlmZmZ ddlm	Z	m
Z
 ddlmZ ddlmZ dd	lmZ dd
lmZ ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ  e j:                  e j<                  d      Z G d d      Z  G d d      Z!y)    N)JSONDecodeError   )ApiError)AsyncClientWrapperSyncClientWrapper)AsyncHttpResponseHttpResponse)parse_obj_as)RequestOptions)BadGatewayError)InternalServerError)ServiceUnavailableError)UnprocessableEntityError)ErrorMessage)OllamaConnectionTestResponse)OllamaModel.c                       e Zd ZdefdZedddedej                  e   dej                  e	   de
ej                  e      fd	Zedddedej                  e   dej                  e	   de
e   fd
Zy)RawOllamaClientclient_wrapperc                    || _         y N_client_wrapperselfr   s     x/Users/manta/Documents/Projects/TheRoad-I1/backend/.venv/lib/python3.12/site-packages/opik/rest_api/ollama/raw_client.py__init__zRawOllamaClient.__init__   
    -    Napi_keyrequest_optionsbase_urlr!   r"   returnc                   | j                   j                  j                  dd||dddi|t              }	 d|j                  cxk  rdk  rin nft        j                  t
        j                  t           t        t
        j                  t           |j                         	            }t        ||
      S |j                  dk(  rt        t        |j                        t        j                  t
        j                  t
        j                      t        t
        j                  t
        j                      |j                         	                  |j                  dk(  rt#        t        |j                        t        j                  t
        j                  t
        j                      t        t
        j                  t
        j                      |j                         	                  |j                  dk(  rUt%        t        |j                        t        j                  t&        t        t&        |j                         	                  |j                         }t+        |j                  t        |j                        |      # t(        $ r6 t+        |j                  t        |j                        |j,                        w xY w)a  
        Fetches the list of models available from the Ollama instance. URL may be provided with or without /v1 suffix (e.g., http://localhost:11434 or http://localhost:11434/v1). The /v1 suffix will be automatically removed for model discovery. For actual LLM inference, use the URL with /v1 suffix for OpenAI-compatible endpoints.

        Parameters
        ----------
        base_url : str
            Base URL of the Ollama instance. May include /v1 suffix which will be automatically removed for connection testing. For inference, use the URL with /v1 suffix for OpenAI-compatible endpoints.

        api_key : typing.Optional[str]
            Optional API key for authenticated Ollama instances. If provided, will be sent as Bearer token in Authorization header.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        HttpResponse[typing.List[OllamaModel]]
            Models retrieved successfully
        v1/private/ollama/modelsPOSTr#   r!   content-typeapplication/jsonmethodjsonheadersr"   omit   ,  type_object_responsedata  r.   body    status_coder.   r:   )r   httpx_clientrequestOMITr>   typingcastListr   r
   r-   r	   r   dictr.   OptionalAnyr   r   r   r   r   textr   r#   r!   r"   	_response_data_response_jsons          r   list_modelszRawOllamaClient.list_models   s5   4 ((55==&$"
  2 , > 
	-	ti++1c1KK, $kk+6 ) 0 $YUCC$$+. !2!23

3$"(//&**"=$-NN$4	 	 $$+) !2!23

3$"(//&**"=$-NN$4	 	 $$+- !2!23$$".$-NN$4	 	 '^^-N 9#8#8$yGXGXBY`noo  	ty'<'<d9K\K\F]dmdrdrss	ts   A=J 4F(J ?Kc                   | j                   j                  j                  dd||dddi|t              }	 d|j                  cxk  rdk  rGn nDt        j                  t        t        t        |j                         	            }t        ||
      S |j                  dk(  rt        t        |j                        t        j                  t
        j                  t
        j                     t        t
        j                  t
        j                     |j                         	                  |j                  dk(  rUt!        t        |j                        t        j                  t        t        t        |j                         	                  |j                  dk(  rUt#        t        |j                        t        j                  t$        t        t$        |j                         	                  |j                         }t)        |j                  t        |j                        |      # t&        $ r6 t)        |j                  t        |j                        |j*                        w xY w)a  
        Validates that the provided Ollama URL is reachable. URL may be provided with or without /v1 suffix (e.g., http://localhost:11434 or http://localhost:11434/v1). The /v1 suffix will be automatically removed for connection testing. For inference, use the URL with /v1 suffix.

        Parameters
        ----------
        base_url : str
            Base URL of the Ollama instance. May include /v1 suffix which will be automatically removed for connection testing. For inference, use the URL with /v1 suffix for OpenAI-compatible endpoints.

        api_key : typing.Optional[str]
            Optional API key for authenticated Ollama instances. If provided, will be sent as Bearer token in Authorization header.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        HttpResponse[OllamaConnectionTestResponse]
            Connection test successful
        !v1/private/ollama/test-connectionr'   r(   r)   r*   r+   r0   r1   r2   r5   r8   r9     r<   r=   )r   r?   r@   rA   r>   rB   rC   r   r
   r-   r	   r   rE   r.   rF   rG   r   r   r   r   r   rH   rI   s          r   test_connectionzRawOllamaClient.test_connectionr   s	   4 ((55==/$"
  2 , > 
	-	ti++1c10 : ) 0 $YUCC$$+. !2!23

3$"(//&**"=$-NN$4	 	 $$+% !2!234$">$-NN$4	 	 $$+- !2!23$$".$-NN$4	 	 '^^-N 9#8#8$yGXGXBY`noo  	ty'<'<d9K\K\F]dmdrdrss	ts   AH/ E2H/ /?I.)__name__
__module____qualname__r   r   rA   strrB   rF   r   r	   rD   r   rM   r   rQ    r   r   r   r      s    .*; . )-;?Up Up %	Up
  8Up 
fkk+.	/Upv )-;?Up Up %	Up
  8Up 
2	3Upr   r   c                       e Zd ZdefdZedddedej                  e   dej                  e	   de
ej                  e      fd	Zedddedej                  e   dej                  e	   de
e   fd
Zy)AsyncRawOllamaClientr   c                    || _         y r   r   r   s     r   r   zAsyncRawOllamaClient.__init__   r   r   Nr    r#   r!   r"   r$   c                  K   | j                   j                  j                  dd||dddi|t               d{   }	 d|j                  cxk  rd	k  rin nft        j                  t
        j                  t           t        t
        j                  t           |j                         
            }t        ||      S |j                  dk(  rt        t        |j                        t        j                  t
        j                  t
        j                      t        t
        j                  t
        j                      |j                         
                  |j                  dk(  rt#        t        |j                        t        j                  t
        j                  t
        j                      t        t
        j                  t
        j                      |j                         
                  |j                  dk(  rUt%        t        |j                        t        j                  t&        t        t&        |j                         
                  |j                         }t+        |j                  t        |j                        |      7 W# t(        $ r6 t+        |j                  t        |j                        |j,                        w xY ww)a  
        Fetches the list of models available from the Ollama instance. URL may be provided with or without /v1 suffix (e.g., http://localhost:11434 or http://localhost:11434/v1). The /v1 suffix will be automatically removed for model discovery. For actual LLM inference, use the URL with /v1 suffix for OpenAI-compatible endpoints.

        Parameters
        ----------
        base_url : str
            Base URL of the Ollama instance. May include /v1 suffix which will be automatically removed for connection testing. For inference, use the URL with /v1 suffix for OpenAI-compatible endpoints.

        api_key : typing.Optional[str]
            Optional API key for authenticated Ollama instances. If provided, will be sent as Bearer token in Authorization header.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        AsyncHttpResponse[typing.List[OllamaModel]]
            Models retrieved successfully
        r&   r'   r(   r)   r*   r+   Nr0   r1   r2   r5   r8   r9   r;   r<   r=   )r   r?   r@   rA   r>   rB   rC   rD   r   r
   r-   r   r   rE   r.   rF   rG   r   r   r   r   r   rH   rI   s          r   rM   z AsyncRawOllamaClient.list_models   sE    4 ..;;CC&$"
  2 , D 
 
	-	ti++1c1KK, $kk+6 ) 0 ))%HH$$+. !2!23

3$"(//&**"=$-NN$4	 	 $$+) !2!23

3$"(//&**"=$-NN$4	 	 $$+- !2!23$$".$-NN$4	 	 '^^-N 9#8#8$yGXGXBY`noow
r  	ty'<'<d9K\K\F]dmdrdrss	ts5   8KJK A=J =K>F(J &,K?KKc                   K   | j                   j                  j                  dd||dddi|t               d{   }	 d|j                  cxk  rd	k  rGn nDt        j                  t        t        t        |j                         
            }t        ||      S |j                  dk(  rt        t        |j                        t        j                  t
        j                  t
        j                     t        t
        j                  t
        j                     |j                         
                  |j                  dk(  rUt!        t        |j                        t        j                  t        t        t        |j                         
                  |j                  dk(  rUt#        t        |j                        t        j                  t$        t        t$        |j                         
                  |j                         }t)        |j                  t        |j                        |      7 # t&        $ r6 t)        |j                  t        |j                        |j*                        w xY ww)a  
        Validates that the provided Ollama URL is reachable. URL may be provided with or without /v1 suffix (e.g., http://localhost:11434 or http://localhost:11434/v1). The /v1 suffix will be automatically removed for connection testing. For inference, use the URL with /v1 suffix.

        Parameters
        ----------
        base_url : str
            Base URL of the Ollama instance. May include /v1 suffix which will be automatically removed for connection testing. For inference, use the URL with /v1 suffix for OpenAI-compatible endpoints.

        api_key : typing.Optional[str]
            Optional API key for authenticated Ollama instances. If provided, will be sent as Bearer token in Authorization header.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        AsyncHttpResponse[OllamaConnectionTestResponse]
            Connection test successful
        rO   r'   r(   r)   r*   r+   Nr0   r1   r2   r5   r8   r9   rP   r<   r=   )r   r?   r@   rA   r>   rB   rC   r   r
   r-   r   r   rE   r.   rF   rG   r   r   r   r   r   rH   rI   s          r   rQ   z$AsyncRawOllamaClient.test_connection%  s    4 ..;;CC/$"
  2 , D 
 
	-	ti++1c10 : ) 0 ))%HH$$+. !2!23

3$"(//&**"=$-NN$4	 	 $$+% !2!234$">$-NN$4	 	 $$+- !2!23$$".$-NN$4	 	 '^^-N 9#8#8$yGXGXBY`noow
r  	ty'<'<d9K\K\F]dmdrdrss	ts5   8I>H9I> AH< I>E2H< ,I><?I;;I>)rR   rS   rT   r   r   rA   rU   rB   rF   r   r   rD   r   rM   r   rQ   rV   r   r   rX   rX      s    .*< . )-;?Up Up %	Up
  8Up 
6;;{3	4Upv )-;?Up Up %	Up
  8Up 
7	8Upr   rX   )"rB   json.decoderr   core.api_errorr   core.client_wrapperr   r   core.http_responser   r	   core.pydantic_utilitiesr
   core.request_optionsr   errors.bad_gateway_errorr   errors.internal_server_errorr    errors.service_unavailable_errorr   !errors.unprocessable_entity_errorr   types.error_messager   %types.ollama_connection_test_responser   types.ollama_modelr   rC   rG   rA   r   rX   rV   r   r   <module>ri      sf     ( % G @ 2 1 6 > F H . P , v{{6::s#pp ppfpp ppr   