Coverage for src/chuck_data/llm/client.py: 0%
26 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
1# src/llm_client.py
2"""
3Client for communicating with Databricks-hosted LLM APIs.
4"""
6from openai import OpenAI
7import logging
8from ..config import get_workspace_url
9from ..databricks_auth import get_databricks_token
11logging.getLogger("openai").setLevel(logging.ERROR)
12logging.getLogger("httpx").setLevel(logging.ERROR)
15class LLMClient:
16 def __init__(self):
17 """
18 Initialize the LLM client.
19 Args:
20 client: Optional pre-configured OpenAI client
21 """
22 # Get token from config or environment
23 try:
24 self.databricks_token = get_databricks_token()
25 except Exception as e:
26 logging.error(f"Error getting Databricks token for LLM client: {e}")
27 self.databricks_token = None
29 self.base_url = get_workspace_url()
31 def chat(self, messages, model=None, tools=None, stream=False, tool_choice="auto"):
32 """
33 Send a chat request to the model.
34 Args:
35 messages: List of message objects
36 model: Model to use (default from config)
37 tools: List of tools to provide
38 stream: Whether to stream the response
39 Returns:
40 Response from the API
41 """
42 # Get active model from config if not specified
43 if not model:
44 try:
45 from ..config import get_active_model
47 model = get_active_model()
48 except ImportError:
49 model = None
51 client = OpenAI(
52 api_key=self.databricks_token,
53 base_url=f"{self.base_url}/serving-endpoints",
54 )
55 if tools:
56 response = client.chat.completions.create(
57 model=model,
58 messages=messages,
59 tools=tools,
60 stream=stream,
61 tool_choice=tool_choice,
62 )
63 else:
64 response = client.chat.completions.create(
65 model=model, messages=messages, stream=stream
66 )
67 return response