change to Mistral

TextGen/gemini.py

@@ -1,18 +1,21 @@
-import google.generativeai as genai
-from google.api_core import retry
+from mistralai import Mistral
 import os
 import pathlib
 import textwrap
 import json
 
-GOOGLE_API_KEY = os.environ['GOOGLE_API_KEY']
+MISTRAL_API_KEY = os.environ['MISTRAL_API_KEY']
 
-genai.configure(api_key=GOOGLE_API_KEY)
-
-model = genai.GenerativeModel(model_name='models/gemini-2.5-pro')
+client = Mistral(api_key=MISTRAL_API_KEY)
+model = "mistral-large-latest"
 
 def generate_story(available_items):
-    response = model.generate_content(f"""
+    response = client.chat.complete(
+        model=model,
+        messages=[
+            {
+                "role": "user",
+                "content": f"""
       You are game master in a roguelike game. You should help me create an consise objective for the player to fulfill in the current level.
       Our protagonist, A girl with a red hood is wandering inside a dungeon.
       You are ONLY allowed to only use the following elements and not modify either the map or the enemies:
@@ -22,9 +25,11 @@ def generate_story(available_items):
       Example 2 : Loot two treasure chest then defeat the boss $boss_name to exit.
       A object of type NPC can be talked to but can't move and can't be fought, an ennemy type object can be fought but not talked to. A boss can be fought but not talked to
       To escape, the player will have to pass by the portal.
-      
-      """, request_options={'retry': retry.Retry()})
-    story = response.text
+      """
+            }
+        ]
+    )
+    story = response.choices[0].message.content
     return story
 
 def place_objects(available_items,story,myMap):
@@ -67,12 +72,13 @@ def place_objects(available_items,story,myMap):
     """
 
     print(prompt)
-    response = model.generate_content(
-      prompt,
-      generation_config={'response_mime_type':'application/json'}
+    response = client.chat.complete(
+        model=model,
+        messages=[{"role": "user", "content": prompt}],
+        response_format={"type": "json_object"}
     )
     try:
-        map_placements=json.dumps(json.loads(response.text), indent=4)
+        map_placements=json.dumps(json.loads(response.choices[0].message.content), indent=4)
     except:
         map_placements="error"
     return map_placements

TextGen/router.py

@@ -1,22 +1,16 @@
 import os
 import time
 from io import BytesIO
-from langchain_core.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from fastapi import FastAPI, HTTPException, Query, Request
 from fastapi.responses import StreamingResponse,Response
 from fastapi.middleware.cors import CORSMiddleware
 
-from langchain.chains import LLMChain
-from langchain.prompts import PromptTemplate
+from mistralai import Mistral
 from TextGen.suno import custom_generate_audio, get_audio_information,generate_lyrics
 from TextGen.gemini import generate_story,place_objects,generate_map_markdown
 #from TextGen.diffusion import generate_image
 #from coqui import predict
-from langchain_google_genai import (
-    ChatGoogleGenerativeAI,
-    HarmBlockThreshold,
-    HarmCategory,
-)
 from TextGen import app
 from gradio_client import Client, handle_file
 from typing import List
@@ -24,29 +18,60 @@ from elevenlabs.client import ElevenLabs
 from elevenlabs import Voice, VoiceSettings, stream
 
 
+mistral_client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
+mistral_model = "mistral-large-latest"
+
+bard_tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "PlayLastMusic",
+            "description": "plays the lastest created music",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "Desicion": {"type": "string", "description": "Yes or No"}
+                },
+                "required": ["Desicion"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "CreateLyrics",
+            "description": "create some Lyrics for a new music",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "Desicion": {"type": "string", "description": "Yes or No"}
+                },
+                "required": ["Desicion"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "CreateNewMusic",
+            "description": "create a new music with the Lyrics previously computed",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "Name": {"type": "string", "description": "tags to describe the new music"}
+                },
+                "required": ["Name"]
+            }
+        }
+    }
+]
+
 Eleven_client = ElevenLabs(
   api_key=os.environ["ELEVEN_API_KEY"], # Defaults to ELEVEN_API_KEY
 )
 
 
 Last_message=None
-class PlayLastMusic(BaseModel):
-    '''plays the lastest created music '''
-    Desicion: str = Field(
-        ..., description="Yes or No"
-    )
-
-class CreateLyrics(BaseModel):
-    f'''create some Lyrics for a new music'''
-    Desicion: str = Field(
-        ..., description="Yes or No"
-    )
-
-class CreateNewMusic(BaseModel):
-    f'''create a new music with the Lyrics previously computed'''
-    Name: str = Field(
-        ..., description="tags to describe the new music"
-    )
 
 class SongRequest(BaseModel):
     prompt: str | None  = None
@@ -120,7 +145,7 @@ def generate_text(messages: List[str], npc:str):
         system_prompt=general_npc_prompt+"/n "+main_npc_system_prompts[npc]
     else:
         system_prompt="you're a character in a video game. Play along."
-    print(system_prompt)    
+    print(system_prompt)
     new_messages=[{"role": "user", "content": system_prompt}]
     for index, message in enumerate(messages):
       if index%2==0:
@@ -128,43 +153,33 @@ def generate_text(messages: List[str], npc:str):
       else:
         new_messages.append({"role": "assistant", "content": message})
     print(new_messages)
-    # Initialize the LLM
-    llm = ChatGoogleGenerativeAI(
-        model="gemini-2.5-pro",
-        max_output_tokens=100,
-        temperature=1,
-        safety_settings={
-                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE
-            },
-    )
-    if npc=="bard":
-        llm = llm.bind_tools([PlayLastMusic,CreateNewMusic,CreateLyrics])
 
-    llm_response = llm.invoke(new_messages)
+    kwargs = {
+        "model": mistral_model,
+        "messages": new_messages,
+        "max_tokens": 100,
+        "temperature": 1,
+    }
+    if npc == "bard":
+        kwargs["tools"] = bard_tools
+        kwargs["tool_choice"] = "auto"
+
+    llm_response = mistral_client.chat.complete(**kwargs)
     print(llm_response)
-    return Generate(text=llm_response.content)
+    content = llm_response.choices[0].message.content or ""
+    return Generate(text=content)
 
 
 def inference_model(system_messsage, prompt):
-          
     new_messages=[{"role": "user", "content": system_messsage},{"role": "user", "content": prompt}]
-    llm = ChatGoogleGenerativeAI(
-        model="gemini-2.5-pro",
-        max_output_tokens=100,
+    llm_response = mistral_client.chat.complete(
+        model=mistral_model,
+        messages=new_messages,
+        max_tokens=100,
         temperature=1,
-        safety_settings={
-                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE
-            },
     )
-    llm_response = llm.invoke(new_messages)
     print(llm_response)
-    return Generate(text=llm_response.content)
+    return Generate(text=llm_response.choices[0].message.content)
 
 @app.get("/", tags=["Home"])
 def api_home(request: Request):

requirements.txt

@@ -2,7 +2,7 @@ fastapi==0.99.1
 uvicorn
 requests
 langchain
-langchain_google_genai
+mistralai
 Pillow
 gradio_client
 TTS