fix for mcp output format

app/main.py

@@ -11,9 +11,49 @@ except Exception as e:
     raise
 
 # ---------------------------------------------------------------------
-# Gradio Interface with MCP support
+# MCP - returns raw dictionary format
 # ---------------------------------------------------------------------
-def retriever_interface(query, reports_filter="", sources_filter="", subtype_filter="", year_filter=""):
+def retrieve_mcp(
+    query: str, 
+    reports_filter: str = "", 
+    sources_filter: str = "", 
+    subtype_filter: str = "", 
+    year_filter: str = ""
+) -> list:
+    """
+    Retrieve semantically similar documents from the vector database for MCP clients.
+    
+    Args:
+        query (str): The search query text
+        reports_filter (str): Comma-separated list of specific report filenames (optional)
+        sources_filter (str): Filter by document source type (optional)  
+        subtype_filter (str): Filter by document subtype (optional)
+        year_filter (str): Comma-separated list of years to filter by (optional)
+    
+    Returns:
+        list: List of dictionaries containing document content, metadata, and scores
+    """
+    # Parse filter inputs (convert empty strings to None or lists)
+    reports = [r.strip() for r in reports_filter.split(",") if r.strip()] if reports_filter else []
+    sources = sources_filter.strip() if sources_filter else None
+    subtype = subtype_filter.strip() if subtype_filter else None
+    year = [y.strip() for y in year_filter.split(",") if y.strip()] if year_filter else None
+    
+    # Call retriever function and return raw results
+    results = retrieve_context(
+        query=query,
+        reports=reports,
+        sources=sources,
+        subtype=subtype,
+        year=year
+    )
+    
+    return results
+
+# ---------------------------------------------------------------------
+# UI - returns formatted string 
+# ---------------------------------------------------------------------
+def retrieve_ui(query, reports_filter="", sources_filter="", subtype_filter="", year_filter=""):
     """
     Wrapper function for gradio interface to handle optional filter parameters
     """
@@ -35,59 +75,82 @@ def retriever_interface(query, reports_filter="", sources_filter="", subtype_fil
     # Format results for display
     formatted_results = []
     for i, doc in enumerate(results, 1):
-        metadata_str = ", ".join([f"{k}: {v}" for k, v in doc.get("metadata", {}).items()])
-        formatted_results.append(f"=== Result {i} ===\nContent: {doc['page_content']}\nMetadata: {metadata_str}\n")
+        # Extract content and metadata using the correct keys from HF Spaces API
+        content = doc.get('answer', '')
+        metadata = doc.get('answer_metadata', {})
+        score = doc.get('score', 'N/A')
+        
+        metadata_str = ", ".join([f"{k}: {v}" for k, v in metadata.items()])
+        formatted_results.append(
+            f"=== Result {i} (Score: {score}) ===\n"
+            f"Content: {content}\n"
+            f"Metadata: {metadata_str}\n"
+        )
     
     return "\n".join(formatted_results)
 
-ui = gr.Interface(
-    fn=retriever_interface,
-    inputs=[
-        gr.Textbox(
-            label="Query", 
-            lines=2, 
-            placeholder="Enter your search query here",
-            info="The query to search for in the vector database"
-        ),
-        gr.Textbox(
-            label="Reports Filter (optional)", 
-            lines=1, 
-            placeholder="report1.pdf, report2.pdf",
-            info="Comma-separated list of specific report filenames to search within (leave empty for all)"
-        ),
-        gr.Textbox(
-            label="Sources Filter (optional)", 
-            lines=1, 
-            placeholder="annual_report",
-            info="Filter by document source type (leave empty for all)"
-        ),
-        gr.Textbox(
-            label="Subtype Filter (optional)", 
-            lines=1, 
-            placeholder="financial",
-            info="Filter by document subtype (leave empty for all)"
-        ),
-        gr.Textbox(
-            label="Year Filter (optional)", 
-            lines=1, 
-            placeholder="2023, 2024",
-            info="Comma-separated list of years to filter by (leave empty for all)"
-        ),
-    ],
-    outputs=gr.Textbox(
-        label="Retrieved Context",
-        lines=10,
-        show_copy_button=True
-    ),
-    title="RAG Retrieval Service UI",
-    description="Retrieves semantically similar documents from vector database. Intended for use in RAG pipelines as an MCP server.",
-)
+# Create the Gradio interface with Blocks to support both UI and MCP
+with gr.Blocks() as ui:
+    gr.Markdown("# RAG Retrieval Service UI")
+    gr.Markdown("Retrieves semantically similar documents from vector database. Intended for use in RAG pipelines as an MCP server.")
+    
+    with gr.Row():
+        with gr.Column():
+            query_input = gr.Textbox(
+                label="Query", 
+                lines=2, 
+                placeholder="Enter your search query here",
+                info="The query to search for in the vector database"
+            )
+            reports_input = gr.Textbox(
+                label="Reports Filter (optional)", 
+                lines=1, 
+                placeholder="report1.pdf, report2.pdf",
+                info="Comma-separated list of specific report filenames to search within (leave empty for all)"
+            )
+            sources_input = gr.Textbox(
+                label="Sources Filter (optional)", 
+                lines=1, 
+                placeholder="annual_report",
+                info="Filter by document source type (leave empty for all)"
+            )
+            subtype_input = gr.Textbox(
+                label="Subtype Filter (optional)", 
+                lines=1, 
+                placeholder="financial",
+                info="Filter by document subtype (leave empty for all)"
+            )
+            year_input = gr.Textbox(
+                label="Year Filter (optional)", 
+                lines=1, 
+                placeholder="2023, 2024",
+                info="Comma-separated list of years to filter by (leave empty for all)"
+            )
+            
+            submit_btn = gr.Button("Submit", variant="primary")
+        
+        with gr.Column():
+            output = gr.Textbox(
+                label="Retrieved Context",
+                lines=10,
+                show_copy_button=True
+            )
+    
+    # UI event handler
+    submit_btn.click(
+        fn=retrieve_ui,
+        inputs=[query_input, reports_input, sources_input, subtype_input, year_input],
+        outputs=output
+    )
+    
+    # MCP endpoint
+    gr.api(retrieve_mcp)
 
 # Launch with MCP server enabled
 if __name__ == "__main__":
     ui.launch(
         server_name="0.0.0.0",
-        server_port=7860,  # Different port from reranker
+        server_port=7860, 
         mcp_server=True,
         show_error=True
     )

app/retriever.py

@@ -91,7 +91,6 @@ def create_filter(
 
 def retrieve_context(
     query: str,
-    vectorstore,
     reports: List[str] = None,
     sources: str = None,
     subtype: str = None,