Update TextGen/router.py

TextGen/router.py

@@ -1,6 +1,6 @@
 import os
 import time
-from pydantic import BaseModel
+from langchain_core.pydantic_v1 import BaseModel, Field
 from fastapi import FastAPI, HTTPException, Query, Request
 from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -17,6 +17,23 @@ from TextGen import app
 from gradio_client import Client
 from typing import List
 
+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 Message(BaseModel):
     npc: str | None  = None
     messages: List[str] | None = None
@@ -63,12 +80,19 @@ def generate_text(messages: List[str], npc:str):
     print(new_messages)
     # Initialize the LLM
     llm = ChatGoogleGenerativeAI(
-        model="gemini-pro",
-        max_output_tokens=100,
+        model="gemini-1.5-pro-latest",
+        max_output_tokens=300,
+        temperature=1,
         safety_settings={
-            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-        },
+                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)
     print(llm_response)
     return Generate(text=llm_response.content)