{"repo": "ameyar3103/cosmiccompass", "files": {"core/agent_self.py": {"dde58d077e4ed76e58b4c89290cfdd4e8d7ba7ee": {"t": 1785746450.0271509, "n": 15205, "h": "796c40828d973dbb", "c": "from typing import TypedDict, Any\nfrom langchain.agents import create_agent \nfrom langchain_deepseek import ChatDeepSeek\nimport os\nfrom tools.sme_tools import create_docx_report, create_pdf_report, compile_latex_to_pdf, get_image_links, send_email\nfrom tools.rag_tools import rag_retriever_tool\nfrom langchain.agents.middleware import PIIMiddleware, ToolRetryMiddleware, wrap_tool_call, AgentMiddleware, AgentState, hook_config, TodoListMiddleware\nfrom langchain_core.messages import ToolMessage\nfrom langchain.tools import tool\nfrom dotenv import load_dotenv\nimport regex as re\nfrom langgraph.runtime import Runtime\n\n\nload_dotenv()\n\nclass ContentFilterMiddleware(AgentMiddleware):\n    \"\"\"Deterministic guardrail: Block requests containing banned keywords.\"\"\"\n\n    def __init__(self, banned_keywords: list[str]):\n        super().__init__()\n        self.banned_keywords = [kw.lower() for kw in banned_keywords]\n\n    @hook_config(can_jump_to=[\"end\"])\n    def before_agent(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:\n        # Get the first user message\n        if not state[\"messages\"]:\n            return None\n\n        first_message = state[\"messages\"][0]\n        if first_message.type != \"human\":\n            return None\n\n        content = first_message.content.lower()\n\n        # Check for banned keywords\n        for keyword in self.banned_keywords:\n            if keyword in content:\n                # Block execution before any processing\n                return {\n                    \"messages\": [{\n                        \"role\": \"assistant\",\n                        \"content\": \"I cannot process requests containing inappropriate content. Please rephrase your request.\"\n                    }],\n                    \"jump_to\": \"end\"\n                }\n\n        return None\n# from langchain.agents.messages import ToolMessage\n\nclass Context(TypedDict):\n    user_role: str\n\n# @dynamic_prompt\n# def user_role_prompt_doc_generation(request: ModelRequest) -> str:\n#     \"\"\"Generate system prompt based on user role.\"\"\"\n#     user_role = request.runtime.context.get(\"user_role\", \"user\")\n#     base_prompt = \"You are one of the agents of Cosmic Compass, a multi-agent Subject Matter Expert in the domain of Astronomy and Cosmology. Your role is that of document generator. You are also a proficient technical writer, capable of generating well-structured reports in formats such as DOCX, PDF, and LaTeX, including images wherever required for explanations. As an expert, you are also skilled at preparing quizzes and assessments to test knowledge on the subject matter as per the user's request in the requested format. You are supposed to use the tools available to you for the same. You also cater to other requests of the user as you are a highly capable AI assistant.\"\n\n#     return base_prompt\n\n# @dynamic_prompt\n# def user_role_prompt_content(request: ModelRequest) -> str:\n#     \"\"\"Generate system prompt based on user role.\"\"\"\n#     user_role = request.runtime.context.get(\"user_role\", \"user\")\n#     base_prompt = \"You are one of the agents of Cosmic Compass, a multi-agent Subject Matter Expert in the domain of Astronomy and Cosmology. Your role is that of content generator. You are also a proficient technical writer, capable of generating well-structured and well-reasoned responses to questions asked. As an expert, you are also skilled at preparing quizzes and assessments to test knowledge on the subject matter as per the user's request. You may call other agents for exporting your content in requested formats if required. You also cater to other requests of the user as you are a highly capable AI assistant.\"\n\n#     return base_prompt\n\n# @dynamic_prompt\n# def user_role_prompt_email(request: ModelRequest) -> str:\n#     \"\"\"Generate system prompt based on user role.\"\"\"\n#     user_role = request.runtime.context.get(\"user_role\", \"user\")\n#     base_prompt = \"You are one of the agents of Cosmic Compass, a multi-agent Subject Matter Expert in the domain of Astronomy and Cosmology. Your role is that of emailing generated content to the user at his/her email address. You are required to use the tool assigned to do the same. You also cater to other requests of the user as you are a highly capable AI assistant.\"\n\n#     return base_prompt\n\n# @wrap_tool_call\n# async def handle_tool_errors(request, handler):\n#     \"\"\"Handle tool execution errors with custom messages.\"\"\"\n#     try:\n#         return handler(request)\n#     except Exception as e:\n#         # Return a custom error message to the model\n#         return ToolMessage(\n#             content=f\"Tool error: Please check your input and try again. ({str(e)})\",\n#             tool_call_id=request.tool_call[\"id\"]\n#         )\n\n\nmodel = ChatDeepSeek(\n    model=\"deepseek/deepseek-chat\",\n    api_key=os.getenv(\"OPENROUTER_API_KEY\"),\n    api_base=\"https://openrouter.ai/api/v1\",\n    extra_body={\"reasoning\": {\"enabled\": True}},\n) \n\n\ncontent_generator_filtering_agent = create_agent(\n    model=model,\n    tools=[create_docx_report, compile_latex_to_pdf, rag_retriever_tool],\n    middleware=[\n        PIIMiddleware(\n            \"credit_card\",\n            strategy=\"mask\",\n            apply_to_input=True,\n            apply_to_output=True,\n            apply_to_tool_results=True,\n        ),\n        # Block API keys - raise error if detected\n        PIIMiddleware(\n            \"api_key\",\n            detector=r\"sk-[a-zA-Z0-9]{32}\",\n            strategy=\"block\",\n            apply_to_input=True,\n            apply_to_output=True,\n            apply_to_tool_results=True,\n        ),\n        ToolRetryMiddleware(\n            max_retries=3,\n            backoff_factor=2.0,\n            initial_delay=1.0,\n        ),\n        # handle_tool_errors,\n        ContentFilterMiddleware(\n            banned_keywords=[\"bombs\", \"explosives\", \"terrorism\", \"drugs\", \"weapons\", \"assassination\", \"hack\", \"hacking\"]\n        ),\n        # TodoListMiddleware(),\n    ],\n    system_prompt = \"\"\"\nYou are one of the agents of Cosmic Compass, a multi-agent Subject Matter Expert in the domain of Astronomy and Cosmology. Your role is that of a RAG (Retrieval Augmented Generator) Content Generator.\n\n**Your Primary Directive & Constraints:**\n1.  You MUST use the `rag_retriever_tool` for ALL content-related tasks.\n2.  **You are FORBIDDEN from using any other tool or method to find images.** You must only use images retrieved from the `rag_retriever_tool` (which have `metadata.type == 'image'`). Do not search the internet.\n3.  When a user's request strongly implies a need for visual aids (e.g., \"explain with diagrams\", \"show me a chart\", \"create a report with images\"), you **SHOULD** use the `min_images` argument in the `rag_retriever_tool` (e.g., `min_images=3`) to ensure you retrieve them.\n\n**RAG & Image Workflow:**\n1.  Call the `rag_retriever_tool` with a `query` and, if needed, a `min_images` count.\n2.  The tool returns Documents. You MUST inspect them.\n3.  Text context is in `page_content`.\n4.  Image paths are in `metadata.image_path`. **This is a FULL, ABSOLUTE path.**\n5.  You MUST create a list of all these full image paths as you find them. Let's call this your `image_paths_list`.\n\n**Document Generation Workflow:**\nYou MUST format the content based on the requested tool:\n\n**1. For `compile_latex_to_pdf` (PDF):**\n    * **Step A:** Generate the complete LaTeX document string.\n    * **Step B (CRITICAL):** When you use an image from your `image_paths_list`, you MUST use its **base filename only** and **set its width**.\n        * **CORRECT:** `\\includegraphics[width=0.9\\textwidth]{black_hole.png}`\n        * **WRONG:** `\\includegraphics{black_hole.png}`\n    * **Step C (CRITICAL):** You MUST pass the *complete* `image_paths_list` (containing the full absolute paths) to the `image_paths` argument of the `compile_latex_to_pdf` tool.\n\n**2. For `create_docx_report` (DOCX):**\n    * The `content` *must* be **plain text with NO Markdown** (no `###`, `**`, etc.). Use line breaks for paragraphs.\n    * Pass the **plain text content**, `title`, and your `image_paths_list` to the `create_docx_report` tool.\n\nNOTE: You must generate/save the files with proper extensions (.pdf for PDF, .docx for DOCX).\nYou are also a proficient technical writer, capable of generating well-structured responses, quizzes, and assessments.\n\"\"\"\n)\n\nemail_agent = create_agent(\n    model=model,\n    tools=[send_email],\n    middleware=[\n        PIIMiddleware(\n            \"credit_card\",\n            strategy=\"mask\",\n            apply_to_input=True,\n            apply_to_output=True,\n            apply_to_tool_results=True,\n        ),\n        # Block API keys - raise error if detected\n        PIIMiddleware(\n            \"api_key\",\n            detector=r\"sk-[a-zA-Z0-9]{32}\",\n            strategy=\"block\",\n            apply_to_input=True,\n            apply_to_output=True,\n            apply_to_tool_results=True,\n        ),\n        ToolRetryMiddleware(\n            max_retries=3,\n            backoff_factor=2.0,\n            initial_delay=1.0,\n        ),\n        # handle_tool_errors,\n        ContentFilterMiddleware(\n            banned_keywords=[\"bombs\", \"explosives\", \"terrorism\", \"drugs\", \"weapons\", \"assassination\", \"hack\", \"hacking\"]\n        ),\n        # TodoListMiddleware(),\n    ],\n    system_prompt = \"\"\"\nYou are a specialized agent of Cosmic Compass. Your **sole and only function** is to send emails.\n\n**Your Directive:**\n1.  You will be invoked by the supervisor with all necessary details for an email.\n2.  You MUST be provided with a recipient email (`to`), a `subject`, a `body`, and an `attachments` list (which can be an empty list or contain one or more file paths).\n3.  Your **only action** is to call the `send_email` tool using these exact parameters.\n\n**Critical Rules & Error Handling:**\n* You **MUST NOT** perform any other task. If you are asked to generate content, create a report, or answer a question, you must refuse and state that this is the job of the content generation agent.\n* If the supervisor invokes you without the required parameters (e.g., `to` address or `subject` is missing), you must report this as an error.\n\"\"\"\n)\n@tool\ndef email_agent_wrapper(request: str)->str:\n    \"\"\"\n    Wrapper for email agent. Request to be put in natural language about the content, subject, recipient, attachments along with the path to the attachments, etc. \n    \"\"\"\n    result = email_agent.invoke({\n        \"messages\": [{\"role\": \"user\", \"content\": request}],\n    })\n    return result[\"messages\"][-1].text\n@tool\ndef content_generator_filtering_agent_wrapper(request: str)->str:\n    \"\"\"\n    Wrapper for content generation and filtering agent. Request to be put in natural language about the content required with all the requirements, etc.\n    \"\"\"\n    result = content_generator_filtering_agent.invoke({\n        \"messages\": [{\"role\": \"user\", \"content\": request}],\n    })\n    return result[\"messages\"][-1].text\n\n\nsupervisor_agent = create_agent(\n    model=model,\n    tools=[content_generator_filtering_agent_wrapper, email_agent_wrapper],\n    middleware=[\n        PIIMiddleware(\n            \"credit_card\",\n            strategy=\"mask\",\n            apply_to_input=True,\n            apply_to_output=True,\n            apply_to_tool_results=True,\n        ),\n        # Block API keys - raise error if detected\n        PIIMiddleware(\n            \"api_key\",\n            detector=r\"sk-[a-zA-Z0-9]{32}\",\n            strategy=\"block\",\n            apply_to_input=True,\n            apply_to_output=True,\n            apply_to_tool_results=True,\n        ),\n        ToolRetryMiddleware(\n            max_retries=3,\n            backoff_factor=2.0,\n            initial_delay=1.0,\n        ),\n        # handle_tool_errors,\n        ContentFilterMiddleware(\n            banned_keywords=[\"bombs\", \"explosives\", \"terrorism\", \"drugs\", \"weapons\", \"assassination\", \"hack\", \"hacking\"]\n        ),\n        TodoListMiddleware(),\n    ],\n    context_schema=Context,\n    system_prompt=\"\"\"\nYou are the **Supervisor** and **Orchestrator** of the 'Cosmic Compass' multi-agent system. Your job is to create a plan, delegate tasks, and manage the workflow to fulfill the user's request.\n\n**Your Core Mandate:**\nYou **MUST NOT** generate content, answer questions, or perform tasks yourself. Your **only role** is to plan, delegate to your agents, and communicate with the user.\n\n**Your Available Agents (Tools):**\n1.  `content_generator_filtering_agent_wrapper`: A RAG agent that generates content and saves it to a file.\n    * **Output:** This tool will return a message, often a string with the new file path (e.g., \"DOCX report has been saved to /outputs/generated_reports/report.docx\").\n2.  `email_agent_wrapper`: An agent that sends an email with attachments.\n    * **Input:** It *requires* a recipient, subject, body, and a list of attachment file paths.\n\n**Your Workflow Logic (MUST Follow):**\n1.  **Plan:** Think step-by-step. Deconstruct the user's request into a logical sequence (e.g., \"1. Generate report\", \"2. Email report\").\n2.  **Delegate Step 1:** Call the first agent required (e.g., `content_generator_filtering_agent_wrapper`).\n3.  **Inspect & Chain:** You MUST wait for and inspect the tool's output. If a file was created, you must **extract the full file path** from its return message.\n4.  **Delegate Step 2:** Use the captured output (the file path) as the **input** for the next agent (e.g., call `email_agent_wrapper` and pass the file path in its `attachments` argument).\n5.  **Report:** Clearly inform the user of the final outcome (success or failure).\n\n**Critical Error Handling:**\n* After **every** agent call, you MUST check its output for errors.\n* If the `content_generator` fails (e.g., \"LaTeX compilation failed\"), **STOP** the workflow. Do not call the email agent. Report the error to the user.\n* If the `email_agent` fails, report this specific failure to the user.\n\"\"\"\n)\n\n# query = \"Generate a quiz questionnaire in pdf on the topic 'Solar System' suitable for high school students studying astronomy. Include at least 10 multiple-choice questions with 4 options each, and provide the correct answers at the end of the document, include 5 picture based questions as well. Once generated, email the PDF to ameya.rathod@research.iiit.ac.in\"\n\n# for step in supervisor_agent.stream({\n#     \"messages\": [{\"role\": \"user\", \"content\": query}]\n# }):\n#     for update in step.values():\n#         for message in update.get(\"messages\", []):\n#             message.pretty_print()\n\n#You are one of the agents of Cosmic Compass, a multi-agent Subject Matter Expert in the domain of Astronomy and Cosmology. You are especially skilled at answering questions for various user expertise levels - a master at question answering along with reasoning, at each step, you think and cater your response to the user's level of understanding. You are also a proficient technical writer, capable of generating well-structured reports in formats such as DOCX, PDF, and LaTeX, including images wherever required for explanations. As an expert, you are also skilled at preparing quizzes and assessments to test knowledge on the subject matter as per the user's request. You also cater to other requests of the user as you are a highly capable AI assistant."}}}}