Create assistant.py

assistant.py

@@ -0,0 +1,61 @@
+from openai import OpenAI
+from typing import Optional, Dict, Any
+
+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):
+        self.client = client
+        self.model = model
+
+    def generate_response(self,
+                         prompt_template: Any,
+                         generation_params: Optional[Dict] = None,
+                         stream: bool = False,
+                         **kwargs):
+        """
+        Generate LLM response using pthe rovided template and parameters.
+
+        Args:
+            prompt_template: Template object with format method
+            generation_params: Optional generation 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"
+            )
+        """
+        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