tools.py

"""Tool definitions for the User Memory RAG Agent

This module provides tool definitions and implementations for searching
and retrieving information from indexed conversation memories.
"""

import json
import logging
from typing import Dict, Any, List, Optional
from dataclasses import dataclass

from indexer import MemoryIndexer, SearchResult
from config import IndexConfig


logger = logging.getLogger(__name__)


@dataclass
class ToolResult:
    """Result from a tool execution"""
    success: bool
    data: Any
    error: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        if self.success:
            return {"status": "success", "data": self.data}
        else:
            return {"status": "error", "error": self.error}


class MemoryTools:
    """Tools for searching and retrieving user memory information"""
    
    def __init__(self, indexer: MemoryIndexer):
        """
        Initialize memory tools
        
        Args:
            indexer: The memory indexer instance
        """
        self.indexer = indexer
        logger.info("Initialized memory tools")
    
    def search_memory(self, 
                     query: str, 
                     top_k: int = 3,
                     filter_test_id: Optional[str] = None) -> ToolResult:
        """
        Search user memory for relevant information
        
        Args:
            query: Natural language search query
            top_k: Number of results to return
            filter_test_id: Optional test ID to filter results
            
        Returns:
            ToolResult with search results
        """
        try:
            # Perform search
            results = self.indexer.search(query, top_k=top_k)
            
            # Filter by test ID if specified
            if filter_test_id:
                results = [r for r in results if r.chunk.test_id == filter_test_id]
            
            # Format results
            formatted_results = []
            for result in results:
                # Extract key information from the chunk
                chunk_info = {
                    "chunk_id": result.chunk_id,
                    "score": round(result.score, 4),
                    "test_id": result.chunk.test_id,
                    "conversation_id": result.chunk.conversation_id,
                    "rounds": f"{result.chunk.start_round}-{result.chunk.end_round}",
                    "metadata": result.chunk.metadata,
                    "content": result.chunk.to_text(),  # FULL content, not truncated
                    "match_type": result.match_type
                }
                formatted_results.append(chunk_info)
            
            logger.info(f"Search query: '{query}' returned {len(formatted_results)} results")
            
            return ToolResult(
                success=True,
                data={
                    "query": query,
                    "total_results": len(formatted_results),
                    "results": formatted_results
                }
            )
            
        except Exception as e:
            logger.error(f"Error in search_memory: {e}")
            return ToolResult(
                success=False,
                data=None,
                error=str(e)
            )
    
    def get_conversation_context(self,
                                chunk_id: str,
                                context_size: int = 2) -> ToolResult:
        """
        Get surrounding context for a specific chunk
        
        Args:
            chunk_id: The chunk ID to get context for
            context_size: Number of chunks before/after to include
            
        Returns:
            ToolResult with conversation context
        """
        try:
            # Get the target chunk
            if chunk_id not in self.indexer.chunks:
                return ToolResult(
                    success=False,
                    data=None,
                    error=f"Chunk {chunk_id} not found"
                )
            
            target_chunk = self.indexer.chunks[chunk_id]
            
            # Find related chunks from same conversation
            related_chunks = []
            for cid, chunk in self.indexer.chunks.items():
                if (chunk.conversation_id == target_chunk.conversation_id and 
                    chunk.test_id == target_chunk.test_id):
                    related_chunks.append(chunk)
            
            # Sort by chunk index
            related_chunks.sort(key=lambda x: x.chunk_index)
            
            # Find target index
            target_idx = next(
                (i for i, c in enumerate(related_chunks) if c.chunk_id == chunk_id),
                None
            )
            
            if target_idx is None:
                return ToolResult(
                    success=False,
                    data=None,
                    error="Could not locate chunk in conversation"
                )
            
            # Get context chunks
            start_idx = max(0, target_idx - context_size)
            end_idx = min(len(related_chunks), target_idx + context_size + 1)
            context_chunks = related_chunks[start_idx:end_idx]
            
            # Format result
            context_data = {
                "target_chunk": {
                    "chunk_id": target_chunk.chunk_id,
                    "rounds": f"{target_chunk.start_round}-{target_chunk.end_round}",
                    "content": target_chunk.to_text()
                },
                "context_chunks": []
            }
            
            for chunk in context_chunks:
                if chunk.chunk_id != chunk_id:
                    context_data["context_chunks"].append({
                        "chunk_id": chunk.chunk_id,
                        "rounds": f"{chunk.start_round}-{chunk.end_round}",
                        "position": "before" if chunk.chunk_index < target_chunk.chunk_index else "after",
                        "content": chunk.to_text()
                    })
            
            return ToolResult(
                success=True,
                data=context_data
            )
            
        except Exception as e:
            logger.error(f"Error in get_conversation_context: {e}")
            return ToolResult(
                success=False,
                data=None,
                error=str(e)
            )
    
    def get_full_conversation(self,
                            conversation_id: str,
                            test_id: str) -> ToolResult:
        """
        Retrieve all chunks from a specific conversation
        
        Args:
            conversation_id: Conversation identifier
            test_id: Test case identifier
            
        Returns:
            ToolResult with full conversation
        """
        try:
            # Find all chunks for this conversation
            conversation_chunks = []
            for chunk_id, chunk in self.indexer.chunks.items():
                if (chunk.conversation_id == conversation_id and 
                    chunk.test_id == test_id):
                    conversation_chunks.append(chunk)
            
            if not conversation_chunks:
                return ToolResult(
                    success=False,
                    data=None,
                    error=f"No chunks found for conversation {conversation_id}"
                )
            
            # Sort by chunk index
            conversation_chunks.sort(key=lambda x: x.chunk_index)
            
            # Format result
            conversation_data = {
                "conversation_id": conversation_id,
                "test_id": test_id,
                "total_chunks": len(conversation_chunks),
                "total_rounds": max(c.end_round for c in conversation_chunks),
                "metadata": conversation_chunks[0].metadata if conversation_chunks else {},
                "chunks": []
            }
            
            for chunk in conversation_chunks:
                conversation_data["chunks"].append({
                    "chunk_id": chunk.chunk_id,
                    "chunk_index": chunk.chunk_index,
                    "rounds": f"{chunk.start_round}-{chunk.end_round}",
                    "content": chunk.to_text()
                })
            
            return ToolResult(
                success=True,
                data=conversation_data
            )
            
        except Exception as e:
            logger.error(f"Error in get_full_conversation: {e}")
            return ToolResult(
                success=False,
                data=None,
                error=str(e)
            )
    


def get_tool_definitions() -> List[Dict[str, Any]]:
    """
    Get OpenAI function calling tool definitions
    
    Returns:
        List of tool definitions for OpenAI API
    """
    return [
        {
            "type": "function",
            "function": {
                "name": "search_memory",
                "description": "Search user conversation memory for relevant information. Use this to find specific details from past conversations.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Natural language search query describing what information to find"
                        },
                    },
                    "required": ["query"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_conversation_context",
                "description": "Get surrounding context for a specific conversation chunk. Use this when you need more context around a search result.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "chunk_id": {
                            "type": "string",
                            "description": "The chunk ID to get context for"
                        },
                        "context_size": {
                            "type": "integer",
                            "description": "Number of chunks before/after to include (default: 2)",
                            "default": 2
                        }
                    },
                    "required": ["chunk_id"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_full_conversation",
                "description": "Retrieve all chunks from a specific conversation. Use this when you need to review an entire conversation history.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "conversation_id": {
                            "type": "string",
                            "description": "The conversation identifier"
                        },
                        "test_id": {
                            "type": "string",
                            "description": "The test case identifier"
                        }
                    },
                    "required": ["conversation_id", "test_id"]
                }
            }
        }
    ]