basic agent

agent.py

@@ -0,0 +1,87 @@
+import os
+
+from dotenv import load_dotenv
+from strictjson import strict_json_async
+
+from sarvam import speaker, translator
+
+load_dotenv()
+
+RAG_SYS_PROMPT = None
+RAG_USER_PROMPT = None
+
+AGENT_PROMPT = """You are an AI agent. 
+  You are given three functions - retriever (Retreives information from a database), translator and a speaker (converts text to speech). 
+  The database is a Grade {} {} Textbook. Your task is to assess the user query and determine which function to call. 
+  If the function is to be called, return response as None. If any function is not needed, you can answer to the query yourself. Also identify keywords in the query,
+  """
+
+
+async def llm(system_prompt: str, user_prompt: str) -> str:
+    from groq import AsyncGroq
+
+    client = AsyncGroq(api_key=os.get_env("GROQ_API_KEY"))
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt},
+    ]
+
+    chat_completion = await client.chat.completions.create(
+        messages=messages,
+        model="llama3-70b-8192",
+        temperature=0.3,
+        max_tokens=360,
+        top_p=1,
+        stop=None,
+        stream=False,
+    )
+
+    return chat_completion.choices[0].message.content
+
+
+async def call_agent(user_prompt, grade, subject):
+    system_prompt = AGENT_PROMPT.format(grade, subject)
+
+    result = await strict_json_async(
+        system_prompt=system_prompt,
+        user_prompt=user_prompt,
+        output_format={
+            "function": 'Type of function to call, type: Enum["retriever", "translator", "speaker", "none"]',
+            "keywords": "Array of keywords, type: List[str]",
+            "src_lang": "Identify the language that the user query is in, type: str",
+            "dest_lang": """Identify the target language from the user query if the function is either "translator" or "speaker". If language is not found, return "none", 
+                                    type: Enum["hindi", "bengali", "kannada", "malayalam", "marathi", "odia", "punjabi", "tamil", "telugu", "english", "gujarati", "none"]""",
+            "source": "Identify the sentence that the user wants to translate or speak. Retu 'none', type: Optional[str]",
+            "response": "Your response, type: Optional[str]",
+        },
+        llm=llm,
+    )
+    return result
+
+
+async def function_caller(user_prompt, grade, subject, client):
+    result = call_agent(user_prompt, grade, subject)
+    function = result["function"].lower()
+
+    if function == "none":
+        return result["response"]
+
+    elif function == "retriever":
+        collection = f"{grade}_{subject}"
+
+        data = client.search(collection, user_prompt)
+        data = [i.document for i in data]
+
+        system_prompt = RAG_SYS_PROMPT.format(grade, subject)
+        user_prompt = RAG_USER_PROMPT.format(user_prompt)
+
+        response = await llm(system_prompt, user_prompt)
+
+        return response
+
+    elif function == "translator":
+        return await translator(result["response"], result["src_lang"], result["dest_lang"])
+
+    elif function == "speaker":
+        return await speaker(result["response"], result["dest_lang"])