Update ai_providers.py

ai_providers.py

@@ -1,32 +1,48 @@
 # ai_providers.py
 
-from typing import Dict, Any
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional
+import os
+from dotenv import load_dotenv
 import time
-from groq import Groq
-from config import GROQ_API_KEY, AI_CONFIG
 
-class GroqProvider:
-    def __init__(self):
-        if not GROQ_API_KEY:
-            raise ValueError("GROQ_API_KEY must be set in environment variables")
-        self.client = Groq(api_key=GROQ_API_KEY)
-        self.model = AI_CONFIG["model"]
+# Corrected import statement
+from config import get_ai_config  # Import the function instead of AI_CONFIG
+
+class AIProvider(ABC):
+    @abstractmethod
+    def get_completion(self, prompt: str, **kwargs) -> str:
+        pass
+
+    @abstractmethod
+    def get_config(self) -> Dict[str, Any]:
+        pass
+
+class GroqProvider(AIProvider):
+    def __init__(self, api_key: str, model: str = "deepseek-ai/deepseek-math-7b-instruct"):
+        self.api_key = api_key
+        self.model = model
         
     def get_completion(self, prompt: str, **kwargs) -> str:
-        """Get completion from Groq API with retry logic"""
+        from groq import Groq
+        
+        client = Groq(api_key=self.api_key)
+        
+        # Configure default parameters
         params = {
-            "temperature": kwargs.get("temperature", AI_CONFIG["temperature"]),
-            "max_tokens": kwargs.get("max_tokens", AI_CONFIG["max_tokens"]),
-            "top_p": kwargs.get("top_p", AI_CONFIG["top_p"]),
+            "temperature": kwargs.get("temperature", 0.7),
+            "max_tokens": kwargs.get("max_tokens", 4000),
+            "top_p": kwargs.get("top_p", 1.0),
             "stop": kwargs.get("stop", None)
         }
         
+        # Add retry logic for robustness
         max_retries = 3
         retry_delay = 1
         
         for attempt in range(max_retries):
             try:
-                completion = self.client.chat.completions.create(
+                completion = client.chat.completions.create(
                     model=self.model,
                     messages=[{"role": "user", "content": prompt}],
                     **params
@@ -35,20 +51,88 @@ class GroqProvider:
                 
             except Exception as e:
                 if attempt == max_retries - 1:
-                    raise Exception(f"Failed to get completion from Groq after {max_retries} attempts: {str(e)}")
+                    raise e
                 time.sleep(retry_delay * (attempt + 1))
         
     def get_config(self) -> Dict[str, Any]:
-        """Get Groq configuration"""
         return {
             "config_list": [{
                 "model": self.model,
-                "api_key": GROQ_API_KEY,
-                "temperature": AI_CONFIG["temperature"],
-                "max_tokens": AI_CONFIG["max_tokens"],
-                "top_p": AI_CONFIG["top_p"]
+                "api_key": self.api_key,
+                "temperature": 0.7,
+                "max_tokens": 4000
             }]
         }
 
-# Global instance
-groq_provider = GroqProvider()
+class OpenAIProvider(AIProvider):
+    def __init__(self, api_key: str, model: str = "gpt-4"):
+        self.api_key = api_key
+        self.model = model
+        
+    def get_completion(self, prompt: str, **kwargs) -> str:
+        import openai
+        openai.api_key = self.api_key
+        
+        response = openai.ChatCompletion.create(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            **kwargs
+        )
+        return response.choices[0].message.content
+        
+    def get_config(self) -> Dict[str, Any]:
+        return {
+            "config_list": [{
+                "model": self.model,
+                "api_key": self.api_key,
+                "temperature": 0.7,
+                "max_tokens": 4000
+            }]
+        }
+
+class AnthropicProvider(AIProvider):
+    def __init__(self, api_key: str, model: str = "claude-3-opus-20240229"):
+        self.api_key = api_key
+        self.model = model
+        
+    def get_completion(self, prompt: str, **kwargs) -> str:
+        import anthropic
+        
+        client = anthropic.Client(api_key=self.api_key)
+        response = client.messages.create(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            **kwargs
+        )
+        return response.content[0].text
+        
+    def get_config(self) -> Dict[str, Any]:
+        return {
+            "config_list": [{
+                "model": self.model,
+                "api_key": self.api_key,
+                "temperature": 0.7,
+                "max_tokens": 4000
+            }]
+        }
+
+class AIProviderFactory:
+    @staticmethod
+    def create_provider(provider_type: str, api_key: Optional[str] = None, model: Optional[str] = None) -> AIProvider:
+        if not api_key:
+            load_dotenv()
+            
+        providers = {
+            "openai": (OpenAIProvider, os.getenv("OPENAI_API_KEY"), "gpt-4"),
+            "groq": (GroqProvider, os.getenv("GROQ_API_KEY"), "deepseek-ai/deepseek-r1-distill-llama-70b"),
+            "anthropic": (AnthropicProvider, os.getenv("ANTHROPIC_API_KEY"), "claude-3-opus-20240229")
+        }
+        
+        if provider_type not in providers:
+            raise ValueError(f"Unsupported provider: {provider_type}")
+            
+        ProviderClass, default_api_key, default_model = providers[provider_type]
+        return ProviderClass(
+            api_key=api_key or default_api_key,
+            model=model or default_model
+        )