llms

llms.py

@@ -0,0 +1,60 @@
+from abc import ABC, abstractmethod
+from vertexai.generative_models import GenerativeModel
+from openai import OpenAI
+import json
+from typing import Dict, List, Optional, Union
+
+class LLMProvider(ABC):
+    @abstractmethod
+    def generate(self, prompt: Optional[str] = None, messages: Optional[List[Dict]] = None, **kwargs) -> str:
+        pass
+
+class GeminiProvider(LLMProvider):
+    def generate(self, prompt: Optional[str] = None, messages: Optional[List[Dict]] = None, **kwargs) -> str:
+        model_name = kwargs.get('model', 'gemini-pro')
+        
+        model = GenerativeModel(model_name=model_name)
+        
+        content = prompt
+        if messages and not prompt:
+            content = "\n".join([f"{m['role']}: {m['content']}" for m in messages])
+            
+        response = model.generate_content(content)
+        return response.text
+
+class OpenAIProvider(LLMProvider):
+    def __init__(self, client):
+        self.client = client
+
+    def generate(self, prompt: Optional[str] = None, messages: Optional[List[Dict]] = None, **kwargs) -> str:
+        if prompt and not messages:
+            messages = [{"role": "user", "content": prompt}]
+            
+        completion_kwargs = {
+            "model": kwargs.get("model", "gpt-3.5-turbo"),
+            "messages": messages,
+            "max_tokens": kwargs.get("max_tokens", 1000),
+            "temperature": kwargs.get("temperature", 1.0),
+        }
+        
+        if "response_format" in kwargs:
+            completion_kwargs["response_format"] = kwargs["response_format"]
+            
+        response = self.client.chat.completions.create(**completion_kwargs)
+        return response.choices[0].message.content
+
+class LLMService:
+    def __init__(self, provider: LLMProvider):
+        self.provider = provider
+
+    def chat(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
+        **kwargs
+    ) -> str:
+        try:
+            return self.provider.generate(prompt=prompt, messages=messages, **kwargs)
+        except Exception as e:
+            print(f"LLM API error: {str(e)}")
+            raise