RAG Module API Reference¶

class LLMFactory:
    """Factory for creating LLM instances based on provider configuration."""

    @staticmethod
    def create_llm(
        provider: str | None = None,
        model: str | None = None,
        temperature: float = 0.2,
        max_tokens: int = 500,
    ) -> BaseLanguageModel:
        """Create an LLM instance based on the provider.

        Args:
        ----
            provider: LLM provider (openai, azure, bedrock, anthropic).
                     Defaults to settings.llm_provider
            model: Model name. Defaults to settings.llm_model
            temperature: Sampling temperature
            max_tokens: Maximum tokens in response

        Returns:
        -------
            LangChain BaseLanguageModel instance

        Raises:
        ------
            ValueError: If provider is not supported or required credentials are missing

        """
        provider = provider or settings.llm_provider
        model = model or settings.llm_model

        if provider == LLMProvider.OPENAI:
            return LLMFactory._create_openai(model, temperature, max_tokens)
        elif provider == LLMProvider.AZURE:
            return LLMFactory._create_azure_openai(model, temperature, max_tokens)
        elif provider == LLMProvider.BEDROCK:
            return LLMFactory._create_bedrock(model, temperature, max_tokens)
        elif provider == LLMProvider.ANTHROPIC:
            return LLMFactory._create_anthropic(model, temperature, max_tokens)
        else:
            msg = f"Unsupported LLM provider: {provider}"
            raise ValueError(msg)

    @staticmethod
    def _create_openai(
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> BaseLanguageModel:
        """Create OpenAI LLM instance."""
        from langchain_openai import ChatOpenAI
        from pydantic import SecretStr

        if not settings.openai_api_key:
            msg = "OPENAI_API_KEY is required for OpenAI provider"
            raise ValueError(msg)

        return ChatOpenAI(
            model=model,
            temperature=temperature,
            max_completion_tokens=max_tokens,
            api_key=SecretStr(settings.openai_api_key),
        )

    @staticmethod
    def _create_azure_openai(
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> BaseLanguageModel:
        """Create Azure OpenAI LLM instance."""
        from langchain_openai import AzureChatOpenAI
        from pydantic import SecretStr

        if not settings.azure_openai_api_key or not settings.azure_openai_endpoint:
            msg = "AZURE_OPENAI_API_KEY and AZURE_OPENAI_ENDPOINT are required for Azure provider"
            raise ValueError(msg)

        deployment_name = settings.azure_openai_deployment or model

        return AzureChatOpenAI(
            azure_deployment=deployment_name,
            model=model,
            temperature=temperature,
            # max_tokens parameter removed - causes empty responses with some models
            azure_endpoint=settings.azure_openai_endpoint,
            api_key=SecretStr(settings.azure_openai_api_key),
            api_version=settings.azure_openai_api_version,
        )

    @staticmethod
    def _create_bedrock(
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> BaseLanguageModel:
        """Create AWS Bedrock LLM instance."""
        from langchain_aws import ChatBedrock
        from pydantic import SecretStr

        if not settings.aws_access_key_id or not settings.aws_secret_access_key:
            msg = "AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY are required for Bedrock provider"
            raise ValueError(msg)

        model_id = settings.bedrock_model_id or model

        return ChatBedrock(
            model=model_id,
            model_kwargs={
                "temperature": temperature,
                "max_tokens": max_tokens,
            },
            credentials_profile_name=None,
            aws_access_key_id=SecretStr(settings.aws_access_key_id),
            aws_secret_access_key=SecretStr(settings.aws_secret_access_key),
            region=settings.aws_region,
        )

    @staticmethod
    def _create_anthropic(
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> BaseLanguageModel:
        """Create Anthropic LLM instance."""
        from langchain_anthropic import ChatAnthropic
        from pydantic import SecretStr

        if not settings.anthropic_api_key:
            msg = "ANTHROPIC_API_KEY is required for Anthropic provider"
            raise ValueError(msg)

        return ChatAnthropic(
            model_name=model,
            temperature=temperature,
            max_tokens_to_sample=max_tokens,
            timeout=None,
            stop=None,
            api_key=SecretStr(settings.anthropic_api_key),
        )

class VectorStoreFactory:
    """Factory for creating vector store instances based on configuration."""

    @staticmethod
    def create_vector_store(
        embeddings: Embeddings,
        store_type: str | None = None,
        **kwargs,
    ) -> VectorStoreInterface:
        """Create a vector store instance.

        Args:
            embeddings: Embeddings model to use
            store_type: Type of store ('faiss', 'qdrant', 'chromadb').
                       If None, uses settings.vector_store_type
            **kwargs: Additional arguments for specific store implementations

        Returns:
            VectorStoreInterface: Initialized vector store

        Raises:
            ValueError: If store_type is not supported

        Examples:
            >>> from green_gov_rag.rag.embeddings import ChunkEmbedder
            >>> embeddings = ChunkEmbedder().embedder
            >>> store = VectorStoreFactory.create_vector_store(embeddings)

            >>> # Explicitly choose Qdrant
            >>> store = VectorStoreFactory.create_vector_store(
            ...     embeddings,
            ...     store_type='qdrant',
            ...     url='http://localhost:6333'
            ... )
        """
        store_type = store_type or settings.vector_store_type

        logger.info(f"Creating vector store: {store_type}")

        if store_type == "faiss":
            return VectorStoreFactory._create_faiss_store(embeddings, **kwargs)
        elif store_type == "qdrant":
            return VectorStoreFactory._create_qdrant_store(embeddings, **kwargs)
        elif store_type == "chromadb":
            return VectorStoreFactory._create_chroma_store(embeddings, **kwargs)
        else:
            raise ValueError(
                f"Unsupported vector store type: {store_type}. "
                f"Supported types: faiss, qdrant, chromadb"
            )

    @staticmethod
    def _create_faiss_store(
        embeddings: Embeddings,
        **kwargs,
    ) -> VectorStoreInterface:
        """Create FAISS vector store."""
        from green_gov_rag.rag.stores import FAISSVectorStore

        index_path = kwargs.get("index_path") or settings.vector_store_path

        return FAISSVectorStore(
            embeddings=embeddings,
            index_path=index_path,
            **kwargs,
        )

    @staticmethod
    def _create_qdrant_store(
        embeddings: Embeddings,
        **kwargs,
    ) -> VectorStoreInterface:
        """Create Qdrant vector store."""
        from green_gov_rag.rag.stores import QdrantVectorStore

        url: str = kwargs.pop("url", None) or settings.qdrant_url or ""
        if not url:
            raise ValueError(
                "Qdrant URL not configured. Set QDRANT_URL in environment or pass url parameter."
            )

        api_key: str | None = kwargs.pop("api_key", None) or settings.qdrant_api_key
        collection_name: str = str(
            kwargs.pop("collection_name", settings.collection_name)
        )

        return QdrantVectorStore(
            embeddings=embeddings,
            url=url,
            api_key=api_key,
            collection_name=collection_name,
            **kwargs,  # Any remaining kwargs
        )

    @staticmethod
    def _create_chroma_store(
        embeddings: Embeddings,
        **kwargs,
    ) -> VectorStoreInterface:
        """Create ChromaDB vector store."""
        # TODO: Implement ChromaDB store
        raise NotImplementedError(
            "ChromaDB support coming soon. Use 'faiss' or 'qdrant' for now."
        )

    @staticmethod
    def get_available_stores() -> list[str]:
        """Get list of available vector store types.

        Returns:
            List of supported store types
        """
        available = ["faiss"]

        # Check if Qdrant is available
        try:
            import qdrant_client  # noqa: F401

            available.append("qdrant")
        except ImportError:
            pass

        # Check if ChromaDB is available
        try:
            import chromadb  # noqa: F401

            available.append("chromadb")
        except ImportError:
            pass

        return available

    @staticmethod
    def validate_config(store_type: str | None = None) -> dict:
        """Validate configuration for a vector store type.

        Args:
            store_type: Type to validate, or None for current config

        Returns:
            Dictionary with validation results

        Examples:
            >>> VectorStoreFactory.validate_config('qdrant')
            {
                'valid': True,
                'store_type': 'qdrant',
                'issues': [],
                'config': {'url': 'http://localhost:6333', ...}
            }
        """
        store_type = store_type or settings.vector_store_type
        issues = []
        config: dict[str, str | None] = {}

        if store_type == "faiss":
            config["index_path"] = settings.vector_store_path
            if not settings.vector_store_path:
                issues.append("VECTOR_STORE_PATH not configured")

        elif store_type == "qdrant":
            config["url"] = settings.qdrant_url
            config["api_key"] = "***" if settings.qdrant_api_key else None

            if not settings.qdrant_url:
                issues.append("QDRANT_URL not configured")

            # Check if Qdrant client is installed
            try:
                import qdrant_client  # noqa: F401
            except ImportError:
                issues.append(
                    "qdrant_client not installed. "
                    "Install with: pip install qdrant-client langchain-qdrant"
                )

        elif store_type == "chromadb":
            issues.append("ChromaDB not yet implemented")

        else:
            issues.append(f"Unknown store type: {store_type}")

        return {
            "valid": len(issues) == 0,
            "store_type": store_type,
            "issues": issues,
            "config": config,
        }

class ChunkEmbedder:
    def __init__(self, provider: str = "bedrock", model_name: str | None = None):
        """Initialize embedding generator.

        :param provider: "bedrock" or "huggingface"
        :param model_name: Name of the model to use.
        """
        self.provider = provider.lower()
        if self.provider == "huggingface":
            self.model_name = model_name or settings.embedding_model
            self.embedder: HuggingFaceEmbeddings | OpenAIEmbeddings = (
                HuggingFaceEmbeddings(model_name=self.model_name)
            )
        elif self.provider == "bedrock":
            bedrock_model = model_name or settings.bedrock_model_id
            self.model_name = bedrock_model if bedrock_model else "anthropic.claude-v2"
            self.embedder = OpenAIEmbeddings(model=self.model_name)
        else:
            msg = "provider must be 'bedrock' or 'huggingface'"
            raise ValueError(msg)

    def embed_chunks(
        self, chunks: list[dict], batch_size: int = 100, show_progress: bool = True
    ) -> list[dict]:
        """Generate embeddings for a list of chunk dictionaries using batching.

        :param chunks: List of dicts with at least {"content": str, "metadata": dict}
        :param batch_size: Number of chunks to embed per batch (default: 100)
        :param show_progress: Show progress information (default: True)
        :return: List of dicts with {"content", "metadata", "embedding"}
        """
        embedded_chunks = []

        # Filter out empty chunks
        valid_chunks = [
            chunk
            for chunk in chunks
            if chunk.get("content") and str(chunk.get("content")).strip()
        ]

        if not valid_chunks:
            return []

        total_batches = (len(valid_chunks) + batch_size - 1) // batch_size

        for i in range(0, len(valid_chunks), batch_size):
            batch = valid_chunks[i : i + batch_size]
            batch_num = i // batch_size + 1

            # Extract texts and metadata
            texts = [chunk["content"] for chunk in batch]
            metadatas = [chunk.get("metadata", {}) for chunk in batch]

            # Generate embeddings for entire batch at once
            vectors = self.embedder.embed_documents(texts)

            # Combine results
            for text, metadata, vector in zip(texts, metadatas, vectors):
                embedded_chunks.append(
                    {"content": text, "metadata": metadata, "embedding": vector}
                )

            if show_progress and batch_num % 10 == 0:
                print(
                    f"   Processed batch {batch_num}/{total_batches} ({len(embedded_chunks)} chunks)"
                )

        if show_progress:
            print(
                f"   Completed: {len(embedded_chunks)} chunks embedded in {total_batches} batches"
            )

        return embedded_chunks

class Citation:
    """A citation linking answer text to source document."""

    def __init__(
        self,
        source_number: int,
        document: Document,
        text_snippet: str,
        confidence: float = 1.0,
    ):
        """Initialize citation.

        Args:
        ----
            source_number: Citation number (1, 2, 3, etc.)
            document: Source Document object
            text_snippet: Text excerpt that was cited
            confidence: Confidence score for this citation (0-1)

        """
        self.source_number = source_number
        self.document = document
        self.text_snippet = text_snippet
        self.confidence = confidence
        self.metadata = document.metadata

    def get_deep_link(self) -> str | None:
        """Generate deep link to specific page/section in PDF.

        Returns
        -------
            URL with fragment identifier for PDF page

        """
        source_url = self.metadata.get("source_url")
        if not source_url:
            return None

        # Get page number if available
        page = self.metadata.get("page")
        if page is not None:
            # PDF page fragment (page=N)
            return f"{source_url}#page={page}"

        # Get section anchor if available
        section_id = self.metadata.get("section_id")
        if section_id:
            return f"{source_url}#{section_id}"

        return source_url

    def get_section_path(self) -> str | None:
        """Get hierarchical section path (e.g., 'Section 2.1.3').

        Returns
        -------
            Formatted section path string

        """
        # Check for hierarchical metadata from LayoutPDFReader
        section_path = self.metadata.get("section_path")
        if section_path:
            return section_path

        # Fallback: construct from section_number and section_title
        section_num = self.metadata.get("section_number")
        section_title = self.metadata.get("section_title")

        if section_num and section_title:
            return f"Section {section_num}: {section_title}"
        elif section_num:
            return f"Section {section_num}"
        elif section_title:
            return section_title

        return None

    def format_citation_markdown(self) -> str:
        """Format citation as markdown with link.

        Returns
        -------
            Markdown-formatted citation string

        """
        title = self.metadata.get("title", "Untitled Document")
        deep_link = self.get_deep_link()
        section_path = self.get_section_path()

        # Build citation components
        citation_parts = [f"[{self.source_number}]"]

        if deep_link:
            citation_parts.append(f"[{title}]({deep_link})")
        else:
            citation_parts.append(title)

        # Add section path if available
        if section_path:
            citation_parts.append(f"({section_path})")

        # Add page number if no section path
        elif "page" in self.metadata:
            citation_parts.append(f"(p. {self.metadata['page']})")

        return " ".join(citation_parts)

    def to_dict(self) -> dict[str, Any]:
        """Convert citation to dictionary.

        Returns
        -------
            Dict representation of citation

        """
        return {
            "source_number": self.source_number,
            "title": self.metadata.get("title", "Untitled"),
            "text_snippet": self.text_snippet,
            "confidence": self.confidence,
            "deep_link": self.get_deep_link(),
            "section_path": self.get_section_path(),
            "page": self.metadata.get("page"),
            "source_url": self.metadata.get("source_url"),
            "metadata": self.metadata,
        }

class EnhancedResponse:
    """Enhanced RAG response with inline citations and source attribution."""

    def __init__(self, answer: str, sources: list[Document], query: str):
        """Initialize enhanced response.

        Args:
        ----
            answer: Generated answer text
            sources: List of source Documents used
            query: Original user query

        """
        self.answer = answer
        self.sources = sources
        self.query = query
        self.citations: list[Citation] = []
        self._build_citations()

    def _build_citations(self) -> None:
        """Build citation objects from source documents."""
        for i, doc in enumerate(self.sources, start=1):
            # Create citation with snippet from document
            snippet = (
                doc.page_content[:200] + "..."
                if len(doc.page_content) > 200
                else doc.page_content
            )

            citation = Citation(
                source_number=i,
                document=doc,
                text_snippet=snippet,
                confidence=doc.metadata.get("relevance_score", 1.0),
            )
            self.citations.append(citation)

    def format_answer_with_inline_citations(self) -> str:
        """Format answer with inline citation markers.

        Returns
        -------
            Answer text with inline [1], [2], etc. citations

        """
        # In a production system, this would use NLP to identify
        # which parts of the answer come from which sources
        # For now, add all citations at the end

        answer_with_citations = self.answer

        # Add citation markers if not already present
        if not any(f"[{i}]" in self.answer for i in range(1, len(self.sources) + 1)):
            # Append source indicators
            citation_markers = ", ".join(
                [f"[{i}]" for i in range(1, len(self.sources) + 1)],
            )
            answer_with_citations = f"{self.answer} {citation_markers}"

        return answer_with_citations

    def format_sources_markdown(self) -> str:
        """Format sources as markdown list with deep links.

        Returns
        -------
            Markdown-formatted sources section

        """
        sources_md = ["## Sources\n"]

        for citation in self.citations:
            sources_md.append(citation.format_citation_markdown())
            sources_md.append("")  # Blank line

        return "\n".join(sources_md)

    def format_full_response_markdown(self) -> str:
        """Format complete response with answer and sources.

        Returns
        -------
            Complete markdown response

        """
        parts = [
            f"**Query:** {self.query}\n",
            "## Answer\n",
            self.format_answer_with_inline_citations(),
            "\n",
            self.format_sources_markdown(),
        ]

        return "\n".join(parts)

    def to_dict(self) -> dict[str, Any]:
        """Convert response to dictionary format.

        Returns
        -------
            Dict representation for API/JSON responses

        """
        return {
            "query": self.query,
            "answer": self.answer,
            "answer_with_citations": self.format_answer_with_inline_citations(),
            "citations": [c.to_dict() for c in self.citations],
            "source_count": len(self.sources),
        }