omkar334/agentic_rag

1from dotenv import load_dotenv2from strictjson import strict_json_async3 4from prompts import (5    AGENT_PROMPT,6    EXTRACT_SYS_PROMPT,7    EXTRACT_USER_PROMPT,8    RAG_SYS_PROMPT,9    RAG_USER_PROMPT,10)11from sarvam import speaker, translator12from scraper import extract13 14load_dotenv()15 16 17async def llm(system_prompt: str, user_prompt: str) -> str:18    import os19 20    from groq import AsyncGroq21 22    client = AsyncGroq(api_key=os.getenv("GROQ_API_KEY"))23 24    messages = [25        {"role": "system", "content": system_prompt},26        {"role": "user", "content": user_prompt},27    ]28 29    chat_completion = await client.chat.completions.create(30        messages=messages,31        model="llama3-70b-8192",32        temperature=0.3,33        max_tokens=360,34        top_p=1,35        stop=None,36        stream=False,37    )38 39    return chat_completion.choices[0].message.content40 41 42async def call_agent(user_prompt, collection):43    grade, subject, chapter = collection.split("_")44 45    system_prompt = AGENT_PROMPT.format(grade, subject)46 47    result = await strict_json_async(48        system_prompt=system_prompt,49        user_prompt=user_prompt,50        output_format={51            "function": 'Type of function to call, type: Enum["retriever", "translator", "speaker", "none", "extractor"]',52            "keywords": "Array of keywords, type: List[str]",53            "src_lang": "Identify the language that the user query is in, type: str",54            "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", 55                                    type: Enum["hindi", "bengali", "kannada", "malayalam", "marathi", "odia", "punjabi", "tamil", "telugu", "english", "gujarati", "none"]""",56            "source": "Identify the sentence that the user wants to translate or speak. Else return 'none', type: Optional[str]",57            "url": "Identify if any URL or link is provided in the user query, type: str",58            "response": "Your response, type: Optional[str]",59        },60        llm=llm,61    )62    return result63 64 65async def retriever(user_prompt, collection, client):66    grade, subject, chapter = collection.split("_")67 68    data = client.search(collection, user_prompt)69    data = [i.document for i in data]70 71    system_prompt = RAG_SYS_PROMPT.format(subject, grade)72    user_prompt = RAG_USER_PROMPT.format(data, user_prompt)73 74    return await llm(system_prompt, user_prompt)75 76 77async def extractor(user_prompt, url):78    text = await extract(url)79 80    system_prompt = EXTRACT_SYS_PROMPT.format(url)81    user_prompt = EXTRACT_USER_PROMPT.format(text, user_prompt)82 83    return await llm(system_prompt, user_prompt)84 85 86async def function_caller(user_prompt, collection, client):87    result = await call_agent(user_prompt, collection)88    print(f"Agent log -\n {result} \n\n")89    function = result["function"].lower()90 91    if function == "none":92        return {"text": result["response"]}93 94    elif function == "retriever":95        response = await retriever(user_prompt, collection, client)96        return {"text": response}97 98    elif function == "translator":99        return await translator(result["source"], result["src_lang"], result["dest_lang"])100 101    elif function == "speaker":102        return await speaker(result["source"])103 104    elif function == "extractor":105        response = await extractor(user_prompt, result["url"])106        return {"text": response}107