Update agent.py

agent.py

@@ -3,101 +3,45 @@ 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.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.headers = {
+            "User-Agent": "Mozilla/5.0"
+        }
 
-        self.llm = self._init_model()
+        # Support for multiple model backends
+        self.supported_models = {
+            "huggingface": self.call_huggingface_llm,
+            # You can easily extend this dictionary to support:
+            # "openai": self.call_openai_model,
+            # "lite_llm": self.call_litellm_model,
+            # "custom_server": self.call_custom_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.")
+        self.default_model = "huggingface"
+        self.model_config = {
+            "huggingface": {
+                "api_key": os.getenv("HF_API_TOKEN"),
+                "model_name": "mistralai/Mistral-7B-Instruct-v0.1"
+            }
+        }
+
+    def __call__(self, question: str) -> str:
+        """
+        Main method to process a question. It first tries DuckDuckGo,
+        then scraping, and finally uses a language model if needed.
+        """
+        print(f"Agent received question: {question[:50]}...")
+        answer = self.get_duckduckgo_answer(question)
+        print(f"Agent returning answer: {answer}")
+        return answer.strip()
 
     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"
 
@@ -107,13 +51,20 @@ class DuckDuckGoAgent:
                 data = response.json()
                 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 with DuckDuckGo API: {e}")
+            print(f"Error contacting 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(
@@ -128,13 +79,54 @@ class DuckDuckGoAgent:
                 text = s.get_text().strip()
                 if text:
                     return text[:200]
-            return self.llm.answer(query)
+            print("No useful snippets found, falling back to language model.")
+            return self.call_model_backend(query)
         except Exception as e:
-            print(f"Scraping error: {e}")
-            return self.llm.answer(query)
+            print(f"Error scraping DuckDuckGo: {e}")
+            return self.call_model_backend(query)
 
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question: {question[:50]}...")
-        answer = self.get_duckduckgo_answer(question)
-        print(f"Agent returning answer: {answer}")
-        return answer
+    def call_model_backend(self, prompt: str) -> str:
+        """
+        Dispatch to the selected LLM backend.
+        """
+        if self.default_model in self.supported_models:
+            return self.supported_models[self.default_model](prompt)
+        return "No valid model backend configured."
+
+    def call_huggingface_llm(self, prompt: str) -> str:
+        """
+        Call Hugging Face Inference API as fallback LLM.
+        """
+        config = self.model_config.get("huggingface", {})
+        api_key = config.get("api_key")
+        model = config.get("model_name")
+
+        if not api_key or not model:
+            return "Error: Hugging Face API Token or model not configured."
+
+        url = f"https://api-inference.huggingface.co/models/{model}"
+        headers = {
+            "Authorization": f"Bearer {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."
+        except Exception as e:
+            print(f"Error contacting Hugging Face LLM: {e}")
+            return "Error contacting Hugging Face model."