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

1# src/llm_client.py 

2""" 

3Client for communicating with Databricks-hosted LLM APIs. 

4""" 

5 

6from openai import OpenAI 

7import logging 

8from ..config import get_workspace_url 

9from ..databricks_auth import get_databricks_token 

10 

11logging.getLogger("openai").setLevel(logging.ERROR) 

12logging.getLogger("httpx").setLevel(logging.ERROR) 

13 

14 

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 

28 

29 self.base_url = get_workspace_url() 

30 

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 

46 

47 model = get_active_model() 

48 except ImportError: 

49 model = None 

50 

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