Update assistant.py

assistant.py

@@ -1,61 +1,56 @@
+from typing import Optional, Dict, Any, Union, Generator
+from huggingface_hub import InferenceClient
 from openai import OpenAI
-from typing import Optional, Dict, Any
+from prompt_template import PromptTemplate
 
 class AIAssistant:
-    """
-    A wrapper class for consistent LLM API interactions.
-
-    This class provides:
-    - Unified interface for different LLM providers
-    - Consistent handling of generation parameters
-    - Support for streaming responses
-
-    Attributes:
-        client: Initialized API client (OpenAI, Anthropic, etc.)
-        model: Name of the model to use
-    """
-    def __init__(self, client: OpenAI, model: str):
+    def __init__(self, client: Union[OpenAI, InferenceClient], model: str):
         self.client = client
         self.model = model
 
-    def generate_response(self,
-                         prompt_template: Any,
-                         generation_params: Optional[Dict] = None,
-                         stream: bool = False,
-                         **kwargs):
+    def generate_response(
+        self,
+        prompt_template: PromptTemplate,
+        messages: list[Dict[str, str]],
+        generation_params: Optional[Dict] = None,
+        stream: bool = True,
+    ) -> Generator[str, None, None]:
         """
-        Generate LLM response using pthe rovided template and parameters.
+        Generate LLM response using the provided template and parameters.
 
         Args:
-            prompt_template: Template object with format method
-            generation_params: Optional generation parameters
+            prompt_template: PromptTemplate object containing template and parameters
+            messages: List of message dictionaries with role and content
+            generation_params: Optional generation parameters (overrides template parameters)
             stream: Whether to stream the response
-            **kwargs: Variables for prompt template
 
-        Returns:
-            API response object or streamed response
-
-        Example:
-            assistant.generate_response(
-                prompt_template=template,
-                temperature=0.7,
-                topic="AI safety"
-            )
+        Yields:
+            Streamed response text
         """
-        messages = prompt_template.format(**kwargs)
-        params = generation_params or {}
-
-        completion = self.client.chat.completions.create(
-            model=self.model,
-            messages=messages,
-            stream=stream,
-            **params
-        )
-
-        if stream:
-            for chunk in completion:
-                if chunk.choices[0].delta.content is not None:
-                    print(chunk.choices[0].delta.content, end="")
-            return completion
-
-        return completion
+        params = generation_params or prompt_template.parameters
+
+        # Ensure messages are in correct format
+        formatted_messages = [
+            {"role": msg["role"], "content": str(msg["content"])}
+            for msg in messages
+        ]
+
+        try:
+            completion = self.client.chat.completions.create(
+                model=self.model,
+                messages=formatted_messages,
+                stream=stream,
+                **params
+            )
+
+            if stream:
+                response = ""
+                for chunk in completion:
+                    if chunk.choices[0].delta.content is not None:
+                        response += chunk.choices[0].delta.content
+                        yield response
+            else:
+                return completion.choices[0].message.content
+
+        except Exception as e:
+            yield f"Error generating response: {str(e)}"