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)