feat(api): add support for attestation

magurh · magurh · commit 70ee02713223 · 2025-03-07T19:42:45.000Z
diff --git a/src/flare_ai_rag/api/routes/chat.py b/src/flare_ai_rag/api/routes/chat.py
@@ -2,6 +2,9 @@
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel, Field
 
+from flare_ai_rag.ai import GeminiProvider
+from flare_ai_rag.attestation import Vtpm, VtpmAttestationError
+from flare_ai_rag.prompts import PromptService, SemanticRouterResponse
 from flare_ai_rag.responder import GeminiResponder
 from flare_ai_rag.retriever import QdrantRetriever
 from flare_ai_rag.router import GeminiRouter
@@ -29,26 +32,37 @@ class ChatRouter:
     generation components to handle a conversation in a single endpoint.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         router: APIRouter,
+        ai: GeminiProvider,
         query_router: GeminiRouter,
         retriever: QdrantRetriever,
         responder: GeminiResponder,
+        attestation: Vtpm,
+        prompts: PromptService,
     ) -> None:
         """
         Initialize the ChatRouter.
 
         Args:
             router (APIRouter): FastAPI router to attach endpoints.
-            query_router: Component that classifies the query.
-            retriever: Component that retrieves relevant documents.
-            responder: Component that generates a response.
+            ai (GeminiProvider): AI client used by a simple semantic router
+                to determine if an attestation was requested or if RAG
+                pipeline should be used.
+            query_router: RAG Component that classifies the query.
+            retriever: RAG Component that retrieves relevant documents.
+            responder: RAG Component that generates a response.
+            attestation (Vtpm): Provider for attestation services
+            prompts (PromptService): Service for managing prompts
         """
         self._router = router
+        self.ai = ai
         self.query_router = query_router
         self.retriever = retriever
         self.responder = responder
+        self.attestation = attestation
+        self.prompts = prompts
         self.logger = logger.bind(router="chat")
         self._setup_routes()
 
@@ -65,35 +79,18 @@ async def chat(message: ChatMessage) -> dict[str, str] | None:  # pyright: ignor
             """
             try:
                 self.logger.debug("Received chat message", message=message.message)
-                # Classify the query.
-                classification = self.query_router.route_query(message.message)
-                self.logger.info("Query classified", classification=classification)
-
-                if classification == "ANSWER":
-                    # Retrieve relevant documents.
-                    retrieved_docs = self.retriever.semantic_search(
-                        message.message, top_k=5
-                    )
-                    self.logger.info("Documents retrieved")
-
-                    # Generate the final answer using retrieved context.
-                    answer = self.responder.generate_response(
-                        message.message, retrieved_docs
-                    )
-                    self.logger.info("Response generated", answer=answer)
-                    return {"classification": classification, "response": answer}
-
-                # Map static responses for CLARIFY and REJECT.
-                static_responses = {
-                    "CLARIFY": "Please provide additional context.",
-                    "REJECT": "The query is out of scope.",
-                }
-
-                if classification in static_responses:
-                    return {
-                        "classification": classification,
-                        "response": static_responses[classification],
-                    }
+
+                # If attestation has previously been requested:
+                if self.attestation.attestation_requested:
+                    try:
+                        resp = self.attestation.get_token([message.message])
+                    except VtpmAttestationError as e:
+                        resp = f"The attestation failed with  error:\n{e.args[0]}"
+                    self.attestation.attestation_requested = False
+                    return {"response": resp}
+
+                route = await self.get_semantic_route(message.message)
+                return await self.route_message(route, message.message)
 
             except Exception as e:
                 self.logger.exception("Chat processing failed", error=str(e))
@@ -103,3 +100,120 @@ async def chat(message: ChatMessage) -> dict[str, str] | None:  # pyright: ignor
     def router(self) -> APIRouter:
         """Return the underlying FastAPI router with registered endpoints."""
         return self._router
+
+    async def get_semantic_route(self, message: str) -> SemanticRouterResponse:
+        """
+        Determine the semantic route for a message using AI provider.
+
+        Args:
+            message: Message to route
+
+        Returns:
+            SemanticRouterResponse: Determined route for the message
+        """
+        try:
+            prompt, mime_type, schema = self.prompts.get_formatted_prompt(
+                "semantic_router", user_input=message
+            )
+            route_response = self.ai.generate(
+                prompt=prompt, response_mime_type=mime_type, response_schema=schema
+            )
+            return SemanticRouterResponse(route_response.text)
+        except Exception as e:
+            self.logger.exception("routing_failed", error=str(e))
+            return SemanticRouterResponse.CONVERSATIONAL
+
+    async def route_message(
+        self, route: SemanticRouterResponse, message: str
+    ) -> dict[str, str]:
+        """
+        Route a message to the appropriate handler based on semantic route.
+
+        Args:
+            route: Determined semantic route
+            message: Original message to handle
+
+        Returns:
+            dict[str, str]: Response from the appropriate handler
+        """
+        handlers = {
+            SemanticRouterResponse.RAG_ROUTER: self.handle_rag_pipeline,
+            SemanticRouterResponse.REQUEST_ATTESTATION: self.handle_attestation,
+            SemanticRouterResponse.CONVERSATIONAL: self.handle_conversation,
+        }
+
+        handler = handlers.get(route)
+        if not handler:
+            return {"response": "Unsupported route"}
+
+        return await handler(message)
+
+    async def handle_rag_pipeline(self, _: str) -> dict[str, str]:
+        """
+        Handle attestation requests.
+
+        Args:
+            _: Unused message parameter
+
+        Returns:
+            dict[str, str]: Response containing attestation request
+        """
+        # Step 1. Classify the user query.
+        prompt, mime_type, schema = self.prompts.get_formatted_prompt("rag_router")
+        classification = self.query_router.route_query(
+            prompt=prompt, response_mime_type=mime_type, response_schema=schema
+        )
+        self.logger.info("Query classified", classification=classification)
+
+        if classification == "ANSWER":
+            # Step 2. Retrieve relevant documents.
+            retrieved_docs = self.retriever.semantic_search(_, top_k=5)
+            self.logger.info("Documents retrieved")
+
+            # Step 3. Generate the final answer.
+            answer = self.responder.generate_response(_, retrieved_docs)
+            self.logger.info("Response generated", answer=answer)
+            return {"classification": classification, "response": answer}
+
+        # Map static responses for CLARIFY and REJECT.
+        static_responses = {
+            "CLARIFY": "Please provide additional context.",
+            "REJECT": "The query is out of scope.",
+        }
+
+        if classification in static_responses:
+            return {
+                "classification": classification,
+                "response": static_responses[classification],
+            }
+
+        self.logger.exception("RAG Routing failed")
+        raise ValueError(classification)
+
+    async def handle_attestation(self, _: str) -> dict[str, str]:
+        """
+        Handle attestation requests.
+
+        Args:
+            _: Unused message parameter
+
+        Returns:
+            dict[str, str]: Response containing attestation request
+        """
+        prompt = self.prompts.get_formatted_prompt("request_attestation")[0]
+        request_attestation_response = self.ai.generate(prompt=prompt)
+        self.attestation.attestation_requested = True
+        return {"response": request_attestation_response.text}
+
+    async def handle_conversation(self, message: str) -> dict[str, str]:
+        """
+        Handle general conversation messages.
+
+        Args:
+            message: Message to process
+
+        Returns:
+            dict[str, str]: Response from AI provider
+        """
+        response = self.ai.send_message(message)
+        return {"response": response.text}
diff --git a/src/flare_ai_rag/main.py b/src/flare_ai_rag/main.py
@@ -15,6 +15,8 @@
 
 from flare_ai_rag.ai import GeminiEmbedding, GeminiProvider
 from flare_ai_rag.api import ChatRouter
+from flare_ai_rag.attestation import Vtpm
+from flare_ai_rag.prompts import PromptService
 from flare_ai_rag.responder import GeminiResponder, ResponderConfig
 from flare_ai_rag.retriever import QdrantRetriever, RetrieverConfig, generate_collection
 from flare_ai_rag.router import GeminiRouter, RouterConfig
@@ -24,20 +26,20 @@
 logger = structlog.get_logger(__name__)
 
 
-def setup_router(input_config: dict) -> GeminiRouter:
-    """Initialize the Gemini Provider and the Gemini Router."""
+def setup_router(input_config: dict) -> tuple[GeminiProvider, GeminiRouter]:
+    """Initialize a Gemini Provider for routing."""
     # Setup router config
     router_model_config = input_config["router_model"]
     router_config = RouterConfig.load(router_model_config)
 
     # Setup Gemini client based on Router config
+    # Older version used a system_instruction
     gemini_provider = GeminiProvider(
-        api_key=settings.gemini_api_key,
-        model=router_config.model.model_id,
-        system_instruction=router_config.system_prompt,
+        api_key=settings.gemini_api_key, model=router_config.model.model_id
     )
+    gemini_router = GeminiRouter(client=gemini_provider, config=router_config)
 
-    return GeminiRouter(client=gemini_provider, config=router_config)
+    return gemini_provider, gemini_router
 
 
 def setup_retriever(
@@ -128,8 +130,8 @@ def create_app() -> FastAPI:
     df_docs = pd.read_csv(settings.data_path / "docs.csv", delimiter=",")
     logger.info("Loaded CSV Data.", num_rows=len(df_docs))
 
-    # Set up the RAG components: 1. Gemini Router
-    router_component = setup_router(input_config)
+    # Set up the RAG components: 1. Gemini Provider
+    base_ai, router_component = setup_router(input_config)
 
     # 2a. Set up Qdrant client.
     qdrant_client = setup_qdrant(input_config)
@@ -143,9 +145,12 @@ def create_app() -> FastAPI:
     # Create an APIRouter for chat endpoints and initialize ChatRouter.
     chat_router = ChatRouter(
         router=APIRouter(),
+        ai=base_ai,
         query_router=router_component,
         retriever=retriever_component,
         responder=responder_component,
+        attestation=Vtpm(simulate=settings.simulate_attestation),
+        prompts=PromptService(),
     )
     app.include_router(chat_router.router, prefix="/api/routes/chat", tags=["chat"])
 
diff --git a/src/flare_ai_rag/prompts/service.py b/src/flare_ai_rag/prompts/service.py
@@ -1,19 +1,11 @@
 """
-Prompt Service Module for Flare AI DeFAI
+Prompt Service Module for Flare AI RAG
 
 This module provides a service layer for managing and formatting AI prompts.
 It acts as a wrapper around the PromptLibrary, adding error handling and
 logging capabilities. The service is responsible for retrieving prompts,
 formatting them with provided parameters, and returning the formatted prompts
 along with their associated metadata.
-
-Example:
-    ```python
-    service = PromptService()
-    prompt, mime_type, schema = service.get_formatted_prompt(
-        "token_send", amount="100", address="0x123..."
-    )
-    ```
 """
 
 from typing import Any
@@ -40,17 +32,6 @@ class to provide additional functionality and safety checks.
         library (PromptLibrary): Instance of the prompt library containing all
             prompt templates
         logger (BoundLogger): Structured logger bound with service context
-
-    Example:
-        ```python
-        service = PromptService()
-        try:
-            prompt, mime_type, schema = service.get_formatted_prompt(
-                "token_send", to_address="0x123...", amount=100
-            )
-        except Exception as e:
-            print(f"Failed to format prompt: {e}")
-        ```
     """
 
     def __init__(self) -> None:
@@ -89,19 +70,6 @@ def get_formatted_prompt(
             ValueError: If required format parameters are missing
             Exception: For other formatting or processing errors
 
-        Example:
-            ```python
-            service = PromptService()
-            try:
-                prompt, mime_type, schema = service.get_formatted_prompt(
-                    "token_swap", from_token="ETH", to_token="USDC", amount=1.5
-                )
-            except KeyError:
-                print("Prompt template not found")
-            except ValueError:
-                print("Missing required parameters")
-            ```
-
         Logs:
             - Exceptions during prompt formatting with prompt name and error details
         """
diff --git a/src/flare_ai_rag/router/base.py b/src/flare_ai_rag/router/base.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Any
 
 
 class BaseQueryRouter(ABC):
@@ -7,7 +8,12 @@ class BaseQueryRouter(ABC):
     """
 
     @abstractmethod
-    def route_query(self, query: str) -> str:
+    def route_query(
+        self,
+        prompt: str,
+        response_mime_type: str | None = None,
+        response_schema: Any | None = None,
+    ) -> str:
         """
         Determine the type of the query: ANSWER, CLARIFY, or REJECT.
         """
diff --git a/src/flare_ai_rag/router/router.py b/src/flare_ai_rag/router/router.py
diff --git a/src/flare_ai_rag/settings.py b/src/flare_ai_rag/settings.py
