Update agents/base_agent.py

agents/base_agent.py

@@ -2,7 +2,9 @@
 
 from abc import ABC, abstractmethod
 from typing import Any, Dict
-import ollama
+from groq import Groq
+client = Groq()
+
 
 class BaseAgent(ABC):
     def __init__(self, model_name: str = "llama3.2:3b"):
@@ -10,10 +12,8 @@ class BaseAgent(ABC):
         
     async def get_completion(self, prompt: str) -> str:
         try:
-            response = ollama.chat(model=self.model_name, messages=[
-                {'role': 'user', 'content': prompt}
-            ])
-            return response['message']['content']
+            response = client.chat.completions.create(groq_api_key=os.getenv("GROQ_API_KEY"), messages=[{'role': 'user', 'content': prompt}], model=self.model_name)
+            return response.choices[0].message.content
         except Exception as e:
             raise Exception(f"Error getting completion: {str(e)}")