fix: using kimi

services/model_handler.py

@@ -1,13 +1,20 @@
 import logging
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import streamlit as st
-from agno.agent import Agent
+from agno.models.openai import OpenAILike
 from agno.models.ollama import Ollama
 from agno.tools.arxiv import ArxivTools
 from agno.tools.pubmed import PubmedTools
+import os
+import random
 
-MODEL_PATH = "meta-llama/Llama-3.2-1B"
-
+API_KEYS = [
+    os.getenv("SK1"),
+    os.getenv("SK2"),
+    os.getenv("SK3"),
+    os.getenv("SK4"),
+    os.getenv("SK5")
+]
 class ModelHandler:
     def __init__(self):
         """Initialize the model handler"""
@@ -21,11 +28,14 @@ class ModelHandler:
     
     def _initialize_model(self):
         """Initialize model and tokenizer"""
-        self.model, self.tokenizer = self._load_model()
         self.translator = Agent(
             name="Translator",
             role="You will translate the query to English",
-            model=Ollama(id="llama3.2:1b"),
+            model=OpenAILike(
+                api_key=random.choice(API_KEYS),
+                base_url="https://api.moonshot.cn/v1",
+                id="moonshot-v1-8k"
+            ),
             goal="Translate to English",
             instructions=[
                 "Translate the query to English"
@@ -35,12 +45,17 @@ class ModelHandler:
         self.researcher = Agent(
             name="Researcher",
             role="You are a research scholar who specializes in autism research.",
-            model=Ollama(id="llama3.2:1b"), 
+            model=OpenAILike(
+                api_key=random.choice(API_KEYS),
+                base_url="https://api.moonshot.cn/v1",
+                id="moonshot-v1-8k"
+            ), 
             tools=[ArxivTools(), PubmedTools()],
             instructions=[
-                "You need to understand the context of the question to provide the best answer based on your tools."
-                "Be precise and provide just enough information to be useful",
-                "You must cite the sources used in your answer."
+                "You have ArxivTools and PubmedTools at your disposal. Use them to find relevant papers for the question.",
+                "You need to understand the context of the question to provide the best answer based on your tools.",
+                "Be precise and provide just enough information to be useful.",
+                "You must cite the sources used in your answer.",
                 "You must create an accessible summary.",
                 "The content must be for people without autism knowledge.",
                 "Focus in the main findings of the paper taking in consideration the question.",
@@ -51,7 +66,11 @@ class ModelHandler:
         self.summarizer = Agent(
             name="Summarizer",
             role="You are a specialist in summarizing research papers for people without autism knowledge.",
-            model=Ollama(id="llama3.2:1b"), 
+            model=OpenAILike(
+                api_key=random.choice(API_KEYS),
+                base_url="https://api.moonshot.cn/v1",
+                id="moonshot-v1-8k"
+            ), 
             instructions=[
                 "You must provide just enough information to be useful",
                 "You must cite the sources used in your answer.",
@@ -74,7 +93,11 @@ class ModelHandler:
         self.presenter = Agent(
             name="Presenter",
             role="You are a professional researcher who presents the results of the research.",
-            model=Ollama(id="llama3.2:1b"),
+            model=OpenAILike(
+                api_key=random.choice(API_KEYS),
+                base_url="https://api.moonshot.cn/v1",
+                id="moonshot-v1-8k"
+            ),
             instructions=[
                 "You are multilingual",
                 "You must present the results in a clear and concise manner.",
@@ -90,18 +113,6 @@ class ModelHandler:
         )
         
     
-    @staticmethod
-    @st.cache_resource
-    @st.cache_data
-    def _load_model():
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-            model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
-            return model, tokenizer
-        except Exception as e:
-            logging.error(f"Error loading model: {str(e)}")
-            return None, None
-    
     def generate_answer(self, query: str) -> str:
         try:
             translator = self.translator.run(query, stream=False)