Upload src/labdaps/chat/session.py with huggingface_hub

src/labdaps/chat/session.py

@@ -1,7 +1,7 @@
 import os
 from dataclasses import dataclass
-import google.generativeai as genai
-from src.labdaps.config import GEMINI_MODEL, MAX_HISTORY_MESSAGES
+from groq import Groq
+from src.labdaps.config import GROQ_MODEL, MAX_HISTORY_MESSAGES
 from src.labdaps.retrieval.retriever import RetrievedChunk, retrieve
 from src.labdaps.ingestion.embedder import Embedder
 from src.labdaps.chat.prompts import build_system_prompt
@@ -16,8 +16,7 @@ class ChatResponse:
 class ChatSession:
     def __init__(self, embedder: Embedder):
         self.embedder = embedder
-        genai.configure(api_key=os.environ["GEMINI_API_KEY"])
-        self.model = genai.GenerativeModel(GEMINI_MODEL)
+        self.client = Groq(api_key=os.environ["GROQ_API_KEY"])
         self.history: list[dict] = []
 
     def reset(self):
@@ -28,17 +27,21 @@ class ChatSession:
         system_prompt = build_system_prompt(chunks)
 
         trimmed_history = self.history[-MAX_HISTORY_MESSAGES:]
-        gemini_history = [
-            {"role": "model" if m["role"] == "assistant" else "user", "parts": [m["content"]]}
-            for m in trimmed_history
-        ]
-
-        chat = self.model.start_chat(history=gemini_history)
-        full_prompt = f"{system_prompt}\n\n---\n\nPergunta: {question}"
+        messages = (
+            [{"role": "system", "content": system_prompt}]
+            + trimmed_history
+            + [{"role": "user", "content": question}]
+        )
 
         full_text = ""
-        for chunk in chat.send_message(full_prompt, stream=True):
-            delta = chunk.text or ""
+        stream = self.client.chat.completions.create(
+            model=GROQ_MODEL,
+            messages=messages,
+            stream=True,
+            max_tokens=2048,
+        )
+        for chunk in stream:
+            delta = chunk.choices[0].delta.content or ""
             full_text += delta
             yield delta, chunks