Update agent.py

agent.py

@@ -3,18 +3,101 @@ import requests
 import urllib.parse
 from bs4 import BeautifulSoup
 
+class BaseModel:
+    def answer(self, prompt: str) -> str:
+        raise NotImplementedError("Model must implement the answer method.")
+
+class HfApiModel(BaseModel):
+    def __init__(self, model_name: str, api_token: str):
+        self.model_name = model_name
+        self.api_token = api_token
+
+    def answer(self, prompt: str) -> str:
+        url = f"https://api-inference.huggingface.co/models/{self.model_name}"
+        headers = {
+            "Authorization": f"Bearer {self.api_token}",
+            "Content-Type": "application/json"
+        }
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": 200,
+                "temperature": 0.0
+            }
+        }
+
+        try:
+            response = requests.post(url, headers=headers, json=payload, timeout=30)
+            response.raise_for_status()
+            output = response.json()
+            if isinstance(output, list) and "generated_text" in output[0]:
+                return output[0]["generated_text"].strip()[:200]
+            return "No response generated."
+        except Exception as e:
+            return f"Error from Hugging Face API: {e}"
+
+class LiteLLMModel(BaseModel):
+    def __init__(self, endpoint_url: str):
+        self.url = endpoint_url
+
+    def answer(self, prompt: str) -> str:
+        try:
+            response = requests.post(self.url, json={"input": prompt}, timeout=30)
+            response.raise_for_status()
+            return response.json().get("output", "No output.")
+        except Exception as e:
+            return f"LiteLLM error: {e}"
+
+class OpenAIServerModel(BaseModel):
+    def __init__(self, api_key: str, model: str = "gpt-3.5-turbo"):
+        self.api_key = api_key
+        self.model = model
+
+    def answer(self, prompt: str) -> str:
+        try:
+            response = requests.post(
+                "https://api.openai.com/v1/chat/completions",
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json"
+                },
+                json={
+                    "model": self.model,
+                    "messages": [{"role": "user", "content": prompt}],
+                    "max_tokens": 200,
+                    "temperature": 0.0
+                },
+                timeout=30
+            )
+            response.raise_for_status()
+            data = response.json()
+            return data["choices"][0]["message"]["content"].strip()[:200]
+        except Exception as e:
+            return f"OpenAI error: {e}"
+
 class DuckDuckGoAgent:
     def __init__(self):
         print("DuckDuckGoAgent initialized.")
-        self.headers = {
-            "User-Agent": "Mozilla/5.0"
-        }
+        self.headers = {"User-Agent": "Mozilla/5.0"}
+        self.hf_api_key = os.getenv("HF_API_TOKEN")
+        self.model_type = os.getenv("MODEL_TYPE", "huggingface")
+        self.model_name = os.getenv("MODEL_NAME", "mistralai/Mistral-7B-Instruct-v0.1")
+        self.model_url = os.getenv("MODEL_URL")  # For LiteLLM
+        self.openai_key = os.getenv("OPENAI_API_KEY")
+
+        self.llm = self._init_model()
+
+    def _init_model(self) -> BaseModel:
+        if self.model_type == "openai" and self.openai_key:
+            return OpenAIServerModel(api_key=self.openai_key)
+        elif self.model_type == "litellm" and self.model_url:
+            return LiteLLMModel(endpoint_url=self.model_url)
+        elif self.model_type == "huggingface" and self.hf_api_key:
+            return HfApiModel(model_name=self.model_name, api_token=self.hf_api_key)
+        else:
+            raise ValueError("No valid model configuration found.")
 
     def get_duckduckgo_answer(self, query: str) -> str:
-        """
-        Attempt to get an answer from the DuckDuckGo API.
-        If no abstract text is found, fall back to scraping.
-        """
         search_query = urllib.parse.quote(query)
         url = f"https://api.duckduckgo.com/?q={search_query}&format=json&no_html=1&skip_disambig=1"
 
@@ -22,23 +105,15 @@ class DuckDuckGoAgent:
             response = requests.get(url, timeout=10)
             if response.status_code == 200:
                 data = response.json()
-                # If AbstractText exists and is non-empty, return it
                 if 'AbstractText' in data and data['AbstractText']:
                     return data['AbstractText'][:200]
-                else:
-                    print("No abstract found, falling back to scraping.")
-                    return self.scrape_duckduckgo(query)
-            else:
-                print(f"DuckDuckGo API failed with status: {response.status_code}")
                 return self.scrape_duckduckgo(query)
+            return self.scrape_duckduckgo(query)
         except Exception as e:
-            print(f"Error contacting DuckDuckGo API: {e}")
+            print(f"Error with DuckDuckGo API: {e}")
             return self.scrape_duckduckgo(query)
 
     def scrape_duckduckgo(self, query: str) -> str:
-        """
-        Fallback to scraping DuckDuckGo search results if API fails or no abstract found.
-        """
         print("Using fallback: scraping HTML results.")
         try:
             response = requests.post(
@@ -53,55 +128,12 @@ class DuckDuckGoAgent:
                 text = s.get_text().strip()
                 if text:
                     return text[:200]
-            print("No useful snippets found, falling back to Hugging Face LLM.")
-            return self.call_huggingface_llm(query)
-        except Exception as e:
-            print(f"Error scraping DuckDuckGo: {e}")
-            return self.call_huggingface_llm(query)
-
-    def call_huggingface_llm(self, prompt: str) -> str:
-        """
-        Fallback to Hugging Face LLM if DuckDuckGo API and scraping both fail.
-        """
-        hf_api_key = os.getenv("HF_API_TOKEN")
-        model = "mistralai/Mistral-7B-Instruct-v0.1"
-
-        if not hf_api_key:
-            return "Error: Hugging Face API Token is not configured."
-
-        url = f"https://api-inference.huggingface.co/models/{model}"
-        headers = {
-            "Authorization": f"Bearer {hf_api_key}",
-            "Content-Type": "application/json"
-        }
-
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": 200,
-                "temperature": 0.7
-            }
-        }
-
-        try:
-            response = requests.post(url, headers=headers, json=payload, timeout=30)
-            response.raise_for_status()
-            output = response.json()
-            if isinstance(output, list) and "generated_text" in output[0]:
-                return output[0]["generated_text"].strip()[:200]
-            elif isinstance(output, dict) and "error" in output:
-                return f"HF LLM error: {output['error']}"
-            else:
-                return "No response generated from Hugging Face LLM."
+            return self.llm.answer(query)
         except Exception as e:
-            print(f"Error contacting Hugging Face LLM: {e}")
-            return "Error contacting Hugging Face model."
+            print(f"Scraping error: {e}")
+            return self.llm.answer(query)
 
     def __call__(self, question: str) -> str:
-        """
-        Main entry point for the agent to process a question.
-        It will first attempt DuckDuckGo, then fall back to scraping or Hugging Face LLM.
-        """
         print(f"Agent received question: {question[:50]}...")
         answer = self.get_duckduckgo_answer(question)
         print(f"Agent returning answer: {answer}")