PitchLense MCP

"""
Google Vertex AI RAG MCP tool.

This tool provides integration with Google Vertex AI's RAG (Retrieval-Augmented Generation)
capabilities for enhanced document search and question answering.

Environment variables:
    GOOGLE_CLOUD_PROJECT: Google Cloud project ID (required)
    GOOGLE_APPLICATION_CREDENTIALS: Path to service account key file (required)
    VERTEX_AI_LOCATION: Vertex AI location/region (default: us-central1)
"""

import os
from typing import Any, Dict, List, Optional
import json

# Conditional imports for Google Cloud dependencies
try:
    from google.cloud import aiplatform
    from google.cloud import discoveryengine_v1beta as discoveryengine
    GOOGLE_CLOUD_AVAILABLE = True
except ImportError:
    GOOGLE_CLOUD_AVAILABLE = False

from ..core.base import BaseMCPTool


class VertexAIRAGMCPTool(BaseMCPTool):
    """Google Vertex AI RAG tool for document search and Q&A."""

    def __init__(self):
        """Initialize the Vertex AI RAG tool."""
        super().__init__()
        
        if not GOOGLE_CLOUD_AVAILABLE:
            raise ImportError(
                "Google Cloud dependencies not available. Install with: "
                "pip install google-cloud-aiplatform google-cloud-discoveryengine"
            )
        
        self.project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
        self.location = os.getenv("VERTEX_AI_LOCATION", "us-central1")
        
        if not self.project_id:
            raise ValueError("GOOGLE_CLOUD_PROJECT environment variable is required")
        
        # Initialize Vertex AI
        aiplatform.init(project=self.project_id, location=self.location)
        
        # Initialize Discovery Engine client
        self.discovery_client = discoveryengine.SearchServiceClient()

    def search_documents(
        self,
        query: str,
        search_engine_id: str,
        data_store_id: str,
        num_results: int = 10,
        filter_expression: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Search documents using Vertex AI Discovery Engine.

        Args:
            query: Search query string
            search_engine_id: ID of the search engine
            data_store_id: ID of the data store
            num_results: Maximum number of results to return
            filter_expression: Optional filter expression for results

        Returns:
            Dictionary containing search results and metadata
        """
        try:
            # Construct the search request
            search_request = discoveryengine.SearchRequest(
                serving_config=f"projects/{self.project_id}/locations/{self.location}/dataStores/{data_store_id}/servingConfigs/{search_engine_id}",
                query=query,
                page_size=num_results,
                filter=filter_expression,
            )

            # Execute the search
            response = self.discovery_client.search(search_request)

            # Process results
            results = []
            for result in response.results:
                doc_data = result.document.derived_struct_data
                results.append({
                    "title": doc_data.get("title", ""),
                    "content": doc_data.get("content", ""),
                    "url": doc_data.get("link", ""),
                    "score": result.relevance_score,
                    "metadata": dict(doc_data)
                })

            return {
                "results": results,
                "total_count": len(results),
                "query": query,
                "error": None
            }

        except Exception as e:
            return {
                "results": [],
                "total_count": 0,
                "query": query,
                "error": str(e)
            }

    def rag_query(
        self,
        question: str,
        search_engine_id: str,
        data_store_id: str,
        context_documents: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        """
        Perform RAG (Retrieval-Augmented Generation) query.

        Args:
            question: Question to answer
            search_engine_id: ID of the search engine
            data_store_id: ID of the data store
            context_documents: Optional list of document IDs to focus on

        Returns:
            Dictionary containing answer and supporting documents
        """
        try:
            # First, search for relevant documents
            search_results = self.search_documents(
                query=question,
                search_engine_id=search_engine_id,
                data_store_id=data_store_id,
                num_results=5
            )

            if search_results["error"]:
                return {
                    "answer": "",
                    "supporting_documents": [],
                    "confidence": 0.0,
                    "error": search_results["error"]
                }

            # Extract context from search results
            context_parts = []
            supporting_docs = []
            
            for result in search_results["results"]:
                context_parts.append(result["content"])
                supporting_docs.append({
                    "title": result["title"],
                    "url": result["url"],
                    "relevance_score": result["score"]
                })

            # Combine context for RAG
            context = "\n\n".join(context_parts)
            
            # Use Vertex AI to generate answer based on retrieved context
            answer_prompt = f"""
            Based on the following context documents, please answer the question: {question}
            
            Context:
            {context}
            
            Please provide a comprehensive answer based on the context. If the context doesn't contain enough information to answer the question, please state that clearly.
            """

            # Initialize Vertex AI model for text generation
            from vertexai.language_models import TextGenerationModel
            
            model = TextGenerationModel.from_pretrained("text-bison@001")
            response = model.predict(
                prompt=answer_prompt,
                max_output_tokens=1024,
                temperature=0.1
            )

            return {
                "answer": response.text,
                "supporting_documents": supporting_docs,
                "confidence": 0.8,  # Placeholder confidence score
                "error": None
            }

        except Exception as e:
            return {
                "answer": "",
                "supporting_documents": [],
                "confidence": 0.0,
                "error": str(e)
            }

    def create_document_index(
        self,
        documents: List[Dict[str, Any]],
        data_store_id: str,
        batch_size: int = 100
    ) -> Dict[str, Any]:
        """
        Create or update document index in Vertex AI Discovery Engine.

        Args:
            documents: List of documents to index
            data_store_id: ID of the data store
            batch_size: Number of documents to process in each batch

        Returns:
            Dictionary containing indexing results
        """
        try:
            from google.cloud import discoveryengine_v1beta as discoveryengine
            
            # Initialize document service client
            document_client = discoveryengine.DocumentServiceClient()
            
            indexed_count = 0
            errors = []
            
            # Process documents in batches
            for i in range(0, len(documents), batch_size):
                batch = documents[i:i + batch_size]
                
                for doc in batch:
                    try:
                        # Create document
                        document = discoveryengine.Document(
                            id=doc.get("id", ""),
                            struct_data=doc.get("content", {}),
                            content=discoveryengine.Document.Content(
                                mime_type="application/json",
                                uri=doc.get("uri", "")
                            )
                        )
                        
                        # Create document request
                        request = discoveryengine.CreateDocumentRequest(
                            parent=f"projects/{self.project_id}/locations/{self.location}/dataStores/{data_store_id}/branches/default_branch",
                            document=document
                        )
                        
                        # Create the document
                        response = document_client.create_document(request)
                        indexed_count += 1
                        
                    except Exception as e:
                        errors.append(f"Document {doc.get('id', 'unknown')}: {str(e)}")
            
            return {
                "indexed_count": indexed_count,
                "total_documents": len(documents),
                "errors": errors,
                "success": len(errors) == 0
            }
            
        except Exception as e:
            return {
                "indexed_count": 0,
                "total_documents": len(documents),
                "errors": [str(e)],
                "success": False
            }

    def get_data_stores(self) -> Dict[str, Any]:
        """
        List available data stores in the project.

        Returns:
            Dictionary containing list of data stores
        """
        try:
            from google.cloud import discoveryengine_v1beta as discoveryengine
            
            # Initialize data store service client
            data_store_client = discoveryengine.DataStoreServiceClient()
            
            # List data stores
            parent = f"projects/{self.project_id}/locations/{self.location}"
            request = discoveryengine.ListDataStoresRequest(parent=parent)
            
            response = data_store_client.list_data_stores(request)
            
            data_stores = []
            for data_store in response:
                data_stores.append({
                    "id": data_store.name.split("/")[-1],
                    "name": data_store.display_name,
                    "type": data_store.content_config.name,
                    "state": data_store.state.name
                })
            
            return {
                "data_stores": data_stores,
                "error": None
            }
            
        except Exception as e:
            return {
                "data_stores": [],
                "error": str(e)
            }