chat_service.py

import json
from fastapi import HTTPException, status
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from app.agentic.graphs import get_chat_agent_graph
from app.agentic.states.chat_agent_state import ChatAgentState
from app.models.user import User
from app.models.conversation import Conversation
from app.models.conversation_message import ConversationMessage
from app.utils.lang_smith_tracing import trace_service, add_trace_metadata
from app.utils.helpers import call_mcp_tool
from database import db


class ChatService:
    
    def _create_conversation(self, user_id: int) -> Conversation:
        """Create a new conversation for the user."""
        session = db()
        conversation = Conversation(user_id=user_id)
        session.add(conversation)
        session.commit()
        session.refresh(conversation)
        return conversation
    
    def _get_conversation(self, conversation_id: int, user_id: int) -> Conversation:
        """Get conversation by ID and verify ownership."""
        session = db()
        conversation = session.query(Conversation).filter(
            Conversation.id == conversation_id
        ).first()
        
        if not conversation:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="Conversation not found"
            )
        
        if conversation.user_id != user_id:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="You don't have access to this conversation"
            )
        
        return conversation
    
    def _get_conversation_history(self, conversation_id: int, limit: int = 10) -> list[BaseMessage]:
        """Fetch last N messages from conversation history and convert to LangChain messages."""
        session = db()
        messages = (
            session.query(ConversationMessage)
            .filter(ConversationMessage.conversation_id == conversation_id)
            .order_by(ConversationMessage.created_at.desc())
            .limit(limit)
            .all()
        )
        
        # Reverse to get chronological order
        messages = list(reversed(messages))
        
        history: list[BaseMessage] = []
        for msg in messages:
            role = msg.role if isinstance(msg.role, str) else msg.role.value
            if role == "user":
                history.append(HumanMessage(content=msg.content))
            elif role == "assistant":
                history.append(AIMessage(content=msg.content))
        
        return history
    
    def _save_messages(self, conversation_id: int, messages: list[tuple[str, str]]) -> None:
        """Save multiple messages to the conversation in a single insert."""
        session = db()
        message_objects = [
            ConversationMessage(
                conversation_id=conversation_id,
                role=role,
                content=content
            )
            for role, content in messages
        ]
        session.add_all(message_objects)
        session.commit()
    
    async def _fetch_contextual_memory(self, user_input: str, user_id: str, agent_id: str, run_id: str) -> list[str]:
        """Fetch contextual memory for the user input."""
        memory_result = await call_mcp_tool(
            tool_name="get_contextual_memory",
            tool_input={
                "user_input": user_input,
                "filter_attributes": {
                    "user_id": user_id,
                    "agent_id": agent_id,
                    "run_id": run_id
                }
            }
        )
        
        if not memory_result:
            return []
        
        # Extract results from the MCP tool response structure
        # Structure: [{"type": "text", "text": "{\"results\": [...]}"}]
        try:
            text_content = memory_result[0]["text"]
            # text might be a JSON string, parse it
            if isinstance(text_content, str):
                text_content = json.loads(text_content)
            results_data = text_content["results"]
        except (IndexError, KeyError, TypeError, json.JSONDecodeError) as e:
            print(f"Failed to fetch contextual memory: {e}")
            return []
        
        if not results_data:
            return []
        
        # Extract only memory content to reduce token cost
        return [item["memory"] for item in results_data if isinstance(item, dict) and "memory" in item]
  
    async def _add_to_contextual_memory(self, memory_content: list[dict[str, str]], user_id: str, agent_id: str, run_id: str):
        """Add new memory to the contextual memory store."""
        await call_mcp_tool(
            tool_name="add_to_memory",
            tool_input={
                "message": memory_content,
                "attributes": {
                    "user_id": user_id,
                    "agent_id": agent_id,
                    "run_id": run_id,
                    "metadata": {}
                }
            }
        )
    @trace_service("chat_service", operation="handle_chat", tags=["chat", "langgraph"])
    async def handle_chat(self, user: User, user_input: str, conversation_id: int | None = None):
        """Handle chat message using LangGraph with MCP tools."""
        
        # Create new conversation or validate existing one
        message_history: list[BaseMessage] = []
        if conversation_id is None:
            conversation = self._create_conversation(user.id)
            conv_id = conversation.id
        else:
            # Verify ownership - raises 404/403 if invalid
            conversation = self._get_conversation(conversation_id, user.id)
            conv_id = conversation.id
            # Fetch existing conversation history
            message_history = self._get_conversation_history(conv_id)
        
        user_id = str(user.id)
        agent_id = "app1"  # In production, get this from config or context
        run_id = str(conv_id)
        
        # Fetch contextual memory before invoking graph
        memories = await self._fetch_contextual_memory(user_input, user_id, agent_id, run_id)
                
        # Add metadata to trace
        add_trace_metadata({
            "conversation_id": conv_id,
            "input_length": len(user_input)
        })
        
        initial_state: ChatAgentState = {
            "user_input": user_input,
            "contextual_memory": memories,
            "conversation_history": message_history,  # Last 10 messages from DB
            "messages": [],  # Will be built by chat_agent
            "attributes": {
                "user_id": user_id,
                "agent_id": agent_id,             
                "run_id": run_id,
                "metadata": {}
            },
            "response": None
        }
        
        # Get the compiled graph (cached after first call)
        graph = await get_chat_agent_graph()
        
        # Invoke the graph asynchronously
        graph_response = await graph.ainvoke(initial_state)
        
        response_content = graph_response.get("response", None)
        
        # Save messages to database (bulk insert)
        self._save_messages(conv_id, [
            ("user", user_input),
            ("assistant", response_content or "")
        ])
        
        # Add to contextual memory
        memory_content = [
            {"role": "user", "content": user_input},
            {"role": "assistant", "content": response_content}
        ]
        
        await self._add_to_contextual_memory(memory_content, user_id, agent_id, run_id)
        
        return {
            "conversation_id": conv_id,
            "response": response_content if response_content else "Sorry, I couldn't generate a response."
        }