odootrain/api/odoo_rag_agent.py

"""
ActiveBlue AI Agent — Odoo 18 RAG Specialist
=============================================
Drop-in specialist agent for the ActiveBlue AI system.
Implements the PeerBus interface defined in ACTIVEBLUE_AI_SPEC.md.

Usage:
    from api.odoo_rag_agent import OdooRagAgent

    agent = OdooRagAgent(rag_url="http://localhost:8000")
    result = await agent.ask("How do I run a payroll batch?")
    print(result["answer"])
"""

import json
import httpx
import logging
from typing import AsyncIterator

log = logging.getLogger(__name__)


class OdooRagAgent:
    name        = "odoo18_rag"
    description = "Answers Odoo 18 questions using RAG over official documentation"
    capabilities = [
        "odoo_how_to",
        "odoo_configuration",
        "odoo_troubleshooting",
        "odoo_workflow",
    ]
    privacy_mode = "local"   # uses local Ollama — HIPAA safe

    def __init__(
        self,
        rag_url: str = "http://localhost:8000",
        timeout: int = 120,
        default_model: str | None = None,
    ):
        self.rag_url       = rag_url.rstrip("/")
        self.timeout       = timeout
        self.default_model = default_model

    async def ask(
        self,
        question: str,
        module: str | None = None,
        top_k: int = 6,
        temperature: float = 0.3,
    ) -> dict:
        payload = {"question": question, "top_k": top_k, "temperature": temperature}
        if module:
            payload["module"] = module
        if self.default_model:
            payload["model"] = self.default_model

        async with httpx.AsyncClient(timeout=self.timeout) as client:
            resp = await client.post(f"{self.rag_url}/ask", json=payload)
            resp.raise_for_status()
            return resp.json()

    async def ask_stream(
        self,
        question: str,
        module: str | None = None,
        top_k: int = 6,
        temperature: float = 0.3,
    ) -> AsyncIterator[str]:
        payload = {"question": question, "top_k": top_k, "temperature": temperature}
        if module:
            payload["module"] = module

        async with httpx.AsyncClient(timeout=self.timeout) as client:
            async with client.stream("POST", f"{self.rag_url}/ask/stream", json=payload) as resp:
                async for line in resp.aiter_lines():
                    if line.startswith("data: "):
                        data_str = line[6:]
                        if data_str == "[DONE]":
                            break
                        try:
                            data = json.loads(data_str)
                            if data.get("type") == "token":
                                yield data["content"]
                            elif data.get("type") == "sources":
                                yield json.dumps(data)
                        except json.JSONDecodeError:
                            continue

    async def handle_peer_message(self, message: dict) -> dict:
        """PeerBus message handler for the ActiveBlue Master AI."""
        action  = message.get("action")
        payload = message.get("payload", {})
        req_id  = message.get("request_id")

        if action == "ask":
            result = await self.ask(
                question    = payload.get("question", ""),
                module      = payload.get("module"),
                top_k       = payload.get("top_k", 6),
                temperature = payload.get("temperature", 0.3),
            )
            return {"request_id": req_id, "agent": self.name, "status": "ok", "result": result}

        elif action == "capabilities":
            return {
                "request_id":  req_id,
                "agent":       self.name,
                "capabilities": self.capabilities,
                "description": self.description,
                "privacy_mode": self.privacy_mode,
            }

        elif action == "health":
            return await self.health()

        return {"request_id": req_id, "agent": self.name, "status": "error", "error": f"Unknown action: {action}"}

    async def health(self) -> dict:
        try:
            async with httpx.AsyncClient(timeout=5) as client:
                resp = await client.get(f"{self.rag_url}/health")
                return {"agent": self.name, "status": "ok", "rag": resp.json()}
        except Exception as e:
            return {"agent": self.name, "status": "error", "error": str(e)}

    # ── Module convenience wrappers ───────────────────────────────────────────

    async def ask_accounting(self, question: str) -> dict:
        return await self.ask(question, module="accounting")

    async def ask_payroll(self, question: str) -> dict:
        return await self.ask(question, module="payroll")

    async def ask_inventory(self, question: str) -> dict:
        return await self.ask(question, module="inventory")

    async def ask_crm(self, question: str) -> dict:
        return await self.ask(question, module="crm")

    async def ask_hr(self, question: str) -> dict:
        return await self.ask(question, module="employees")

    async def ask_manufacturing(self, question: str) -> dict:
        return await self.ask(question, module="manufacturing")

    async def ask_helpdesk(self, question: str) -> dict:
        return await self.ask(question, module="helpdesk")