endpoints and models

agent.py

@@ -32,7 +32,9 @@ async def llm(system_prompt: str, user_prompt: str) -> str:
     return chat_completion.choices[0].message.content
 
 
-async def call_agent(user_prompt, grade, subject):
+async def call_agent(user_prompt, collection):
+    grade, subject, chapter = collection.split("_")
+
     system_prompt = AGENT_PROMPT.format(grade, subject)
 
     result = await strict_json_async(
@@ -52,23 +54,27 @@ async def call_agent(user_prompt, grade, subject):
     return result
 
 
-async def function_caller(user_prompt, collection, client):
+async def retriever(user_prompt, collection, client):
     grade, subject, chapter = collection.split("_")
 
-    result = await call_agent(user_prompt, grade, subject)
+    data = client.search(collection, user_prompt)
+    data = [i.document for i in data]
+
+    system_prompt = RAG_SYS_PROMPT.format(subject, grade)
+    user_prompt = RAG_USER_PROMPT.format(data, user_prompt)
+
+    return await llm(system_prompt, user_prompt)
+
+
+async def function_caller(user_prompt, collection, client):
+    result = await call_agent(user_prompt, collection)
     function = result["function"].lower()
 
     if function == "none":
-        return result["response"]
+        return {"text": result["response"]}
 
     elif function == "retriever":
-        data = client.search(collection, user_prompt)
-        data = [i.document for i in data]
-
-        system_prompt = RAG_SYS_PROMPT.format(subject, grade)
-        user_prompt = RAG_USER_PROMPT.format(data, user_prompt)
-
-        response = await llm(system_prompt, user_prompt)
+        response = await retriever(user_prompt, collection, client)
         return {"text": response}
 
     elif function == "translator":

app.py

@@ -6,8 +6,9 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
-from agent import function_caller
+from agent import function_caller, retriever
 from client import HybridClient
+from sarvam import speaker, translator
 
 app = FastAPI()
 hclient = HybridClient()
@@ -23,17 +24,46 @@ app.add_middleware(
 
 class ChatQuery(BaseModel):
     query: str
-    collection: str
+    grade: str
+    subject: str
+    chapter: str
 
 
-@app.get("/chat")
-async def chat(query: ChatQuery):
-    return await function_caller(query.query, query.collection, hclient)
+class TranslateQuery(BaseModel):
+    text: str
+    src: str
+    dest: str
 
 
-async def gradio_interface(input_text, grade, subject, chapter, history):
+class TTSQuery(BaseModel):
+    text: str
+    src: str
+
+
+@app.get("/agent")
+async def agent(query: ChatQuery):
     collection = f"{grade}_{subject.lower()}_{chapter}"
-    response = await chat(ChatQuery(query=input_text, collection=collection))
+    return await function_caller(query.query, collection, hclient)
+
+
+@app.get("/rag")
+async def rag(query: ChatQuery):
+    collection = f"{grade}_{subject.lower()}_{chapter}"
+    return await retriever(query.query, collection, hclient)
+
+
+@app.get("/translate")
+async def translate(query: TranslateQuery):
+    return await translator(query.text, query.src, query.dest)
+
+
+@app.get("/tts")
+async def tts(query: TTSQuery):
+    return await speaker(query.text, query.src)
+
+
+async def gradio_interface(input_text, grade, subject, chapter, history):
+    response = await agent(ChatQuery(query=input_text, grade=grade, subject=subject, chapter=chapter))
 
     if "text" in response:
         output = response["text"]

sarvam.py

@@ -40,13 +40,13 @@ async def translator(text, src, dest):
                 return {"text": output["translated_text"]}
 
 
-async def speaker(text, dest="hindi"):
+async def speaker(text, src="hindi"):
     async with aiohttp.ClientSession() as session:
         url = "https://api.sarvam.ai/text-to-speech"
 
         payload = {
             "inputs": [text],
-            "target_language_code": code_map[dest],
+            "target_language_code": code_map[src],
             "speaker": "meera",
             "pitch": 0,
             "pace": 1.25,