ts cleanup langserve

app/main.py

@@ -791,109 +791,8 @@ async def chatui_adapter(data):
         logger.error(f"ChatUI error: {str(e)}")
         yield f"Error: {str(e)}"
 
-# # MAIN CHATUI STREAMING ENDPOINT - KEEP THIS (ChatUI likely uses this)
-# async def chatui_stream_with_sources(request: Request):
-#     """ChatUI streaming endpoint that supports sources format"""
-#     try:
-#         body = await request.json()
-#         query = body.get("text", "")
-        
-#         if not query:
-#             async def error_stream():
-#                 yield f"event: error\ndata: {json.dumps({'error': 'No query provided'})}\n\n"
-#             return StreamingResponse(error_stream(), media_type="text/event-stream")
-        
-#         async def event_stream():
-#             try:
-#                 # Process the query through the orchestrator pipeline
-#                 start_time = datetime.now()
-#                 session_id = f"chatui_{start_time.strftime('%Y%m%d_%H%M%S')}"
-                
-#                 initial_state = {
-#                     "query": query,
-#                     "context": "",
-#                     "ingestor_context": "",
-#                     "result": "",
-#                     "sources": [],
-#                     "reports_filter": "",
-#                     "sources_filter": "",
-#                     "subtype_filter": "",
-#                     "year_filter": "",
-#                     "file_content": None,
-#                     "filename": None,
-#                     "file_type": "unknown",
-#                     "workflow_type": "standard",
-#                     "metadata": {
-#                         "session_id": session_id,
-#                         "start_time": start_time.isoformat(),
-#                         "has_file_attachment": False
-#                     }
-#                 }
-                
-#                 # Process non-streaming steps
-#                 state_after_detect = {**initial_state, **detect_file_type_node(initial_state)}
-#                 state_after_ingest = {**state_after_detect, **ingest_node(state_after_detect)}
-#                 workflow_type = route_workflow(state_after_ingest)
-                
-#                 if workflow_type == "geojson_direct":
-#                     final_state = geojson_direct_result_node(state_after_ingest)
-#                     yield f"event: data\ndata: {json.dumps(final_state['result'])}\n\n"
-#                     yield f"event: end\ndata: {{}}\n\n"
-#                 else:
-#                     state_after_retrieve = {**state_after_ingest, **retrieve_node(state_after_ingest)}
-                    
-#                     # Stream generation with sources
-#                     sources_sent = False
-#                     final_sources = None
-#                     partial_count = 0
-                    
 
-                    
-#                     async for partial_state in generate_node_streaming(state_after_retrieve):
-#                         if "result" in partial_state:
-#                             sse_message = f"event: data\ndata: {json.dumps(partial_state['result'])}\n\n"
-#                             yield sse_message
-                        
-#                         # Store sources but don't send immediately
-#                         if "sources" in partial_state:
-#                             final_sources = partial_state["sources"]
-                    
-#                     # Send sources as the last event before end
-#                     if final_sources and not sources_sent:
-#                         sources_data = {"sources": final_sources}
-#                         sse_sources_message = f"event: sources\ndata: {json.dumps(sources_data)}\n\n"
-#                         yield sse_sources_message
-#                         sources_sent = True
-                    
-#                     end_message = f"event: end\ndata: {{}}\n\n"
-#                     yield end_message
-                    
-#             except Exception as e:
-#                 logger.error(f"ChatUI streaming error: {str(e)}")
-#                 yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"
-        
-#         return StreamingResponse(
-#             event_stream(),
-#             media_type="text/event-stream",
-#             headers={
-#                 "Cache-Control": "no-cache",
-#                 "Connection": "keep-alive",
-#                 "Access-Control-Allow-Origin": "*",
-#                 "Access-Control-Allow-Headers": "*",
-#             }
-#         )
-        
-#     except Exception as e:
-#         logger.error(f"ChatUI request parsing failed: {str(e)}")
-#         async def error_stream():
-#             yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"
-        
-#         return StreamingResponse(
-#             error_stream(),
-#             media_type="text/event-stream"
-#         )
-
-# TODO: PROBABLY REMOVE - USED BY LANGSERVE ENDPOINT
+# USED BY LANGSERVE ENDPOINT
 def process_query_langserve(input_data: ChatFedInput) -> ChatFedOutput:
     result = process_query_core(
         query=input_data["query"],
@@ -909,37 +808,37 @@ def process_query_langserve(input_data: ChatFedInput) -> ChatFedOutput:
     )
     return ChatFedOutput(result=result["result"], metadata=result["metadata"])
 
-# PROBABLY REMOVE - GRADIO UI
-# def create_gradio_interface():
-#     with gr.Blocks(title="ChatFed Orchestrator") as demo:
-#         gr.Markdown("# ChatFed Orchestrator")
-#         gr.Markdown("Upload documents (PDF/DOCX/GeoJSON) alongside your queries for enhanced context. MCP endpoints available at `/gradio_api/mcp/sse`")
+# GRADIO TEST UI
+def create_gradio_interface():
+    with gr.Blocks(title="ChatFed Orchestrator") as demo:
+        gr.Markdown("# ChatFed Orchestrator")
+        gr.Markdown("Upload documents (PDF/DOCX/GeoJSON) alongside your queries for enhanced context. MCP endpoints available at `/gradio_api/mcp/sse`")
         
-#         with gr.Row():
-#             with gr.Column():
-#                 query_input = gr.Textbox(label="Query", lines=2, placeholder="Enter your question...")
-#                 file_input = gr.File(label="Upload Document (PDF/DOCX/GeoJSON)", file_types=[".pdf", ".docx", ".geojson", ".json"])
+        with gr.Row():
+            with gr.Column():
+                query_input = gr.Textbox(label="Query", lines=2, placeholder="Enter your question...")
+                file_input = gr.File(label="Upload Document (PDF/DOCX/GeoJSON)", file_types=[".pdf", ".docx", ".geojson", ".json"])
                 
-#                 with gr.Accordion("Filters (Optional)", open=False):
-#                     reports_filter_input = gr.Textbox(label="Reports Filter", placeholder="e.g., annual_reports")
-#                     sources_filter_input = gr.Textbox(label="Sources Filter", placeholder="e.g., internal")
-#                     subtype_filter_input = gr.Textbox(label="Subtype Filter", placeholder="e.g., financial")
-#                     year_filter_input = gr.Textbox(label="Year Filter", placeholder="e.g., 2024")
+                with gr.Accordion("Filters (Optional)", open=False):
+                    reports_filter_input = gr.Textbox(label="Reports Filter", placeholder="e.g., annual_reports")
+                    sources_filter_input = gr.Textbox(label="Sources Filter", placeholder="e.g., internal")
+                    subtype_filter_input = gr.Textbox(label="Subtype Filter", placeholder="e.g., financial")
+                    year_filter_input = gr.Textbox(label="Year Filter", placeholder="e.g., 2024")
                 
-#                 submit_btn = gr.Button("Submit", variant="primary")
+                submit_btn = gr.Button("Submit", variant="primary")
             
-#             with gr.Column():
-#                 output = gr.Textbox(label="Response", lines=15, show_copy_button=True)
+            with gr.Column():
+                output = gr.Textbox(label="Response", lines=15, show_copy_button=True)
         
-#         # Use streaming function
-#         submit_btn.click(
-#             fn=process_query_gradio_streaming,
-#             inputs=[query_input, file_input, reports_filter_input, sources_filter_input, 
-#                    subtype_filter_input, year_filter_input],
-#             outputs=output
-#         )
+        # Use streaming function
+        submit_btn.click(
+            fn=process_query_gradio_streaming,
+            inputs=[query_input, file_input, reports_filter_input, sources_filter_input, 
+                   subtype_filter_input, year_filter_input],
+            outputs=output
+        )
     
-#     return demo
+    return demo
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -969,15 +868,10 @@ async def root():
             "chatfed-ui-stream": "/chatfed-ui-stream",  
             "chatfed-with-file": "/chatfed-with-file",
             # "chatfed-with-file-stream": "/chatfed-with-file/stream",
-            # "chatui-stream": "/chatui-stream"
         }
     }
 
-# MAIN CHATUI ENDPOINT - KEEP THIS 
-# @app.post("/chatui-stream")
-# async def chatui_stream_endpoint(request: Request):
-#     """ChatUI compatible streaming endpoint with sources support"""
-#     return await chatui_stream_with_sources(request)
+
 
 # # FILE UPLOAD ADAPTER - KEEP THIS
 async def chatfed_with_file_adapter(
@@ -1113,7 +1007,7 @@ async def chatfed_with_file_stream(
         filename = file.filename
     
     async def generate_sse_stream():
-        """Generate Server-Sent Events format that ChatUI expects"""
+        """Generate Server-Sent Events format for ChatUI"""
         try:
             # Small delay to ensure SSE connection is established
             await asyncio.sleep(0.1)
@@ -1148,7 +1042,7 @@ async def chatfed_with_file_stream(
     
     return StreamingResponse(
         generate_sse_stream(),
-        media_type="text/event-stream",  # This is what ChatUI expects!
+        media_type="text/event-stream", #ChatUI format
         headers={
             "Cache-Control": "no-cache",
             "Connection": "keep-alive",
@@ -1158,13 +1052,13 @@ async def chatfed_with_file_stream(
 
 # TODO: TEST IF CHATUI NEEDS THESE LANGSERVE ENDPOINTS
 # If ChatUI works without these, they can be removed
-add_routes(
-    app,
-    RunnableLambda(process_query_langserve),
-    path="/chatfed",
-    input_type=ChatFedInput,
-    output_type=ChatFedOutput
-)
+# add_routes(
+#     app,
+#     RunnableLambda(process_query_langserve),
+#     path="/chatfed",
+#     input_type=ChatFedInput,
+#     output_type=ChatFedOutput
+# )
 
 add_routes(
     app,
@@ -1177,17 +1071,16 @@ add_routes(
 )
 
 if __name__ == "__main__":
-    # PROBABLY REMOVE - GRADIO INTERFACE 
-    # # Create Gradio interface
-    # demo = create_gradio_interface()
+    # Create Gradio interface
+    demo = create_gradio_interface()
     
-    # # Mount Gradio app to FastAPI
-    # app = gr.mount_gradio_app(app, demo, path="/gradio")
+    # Mount Gradio app to FastAPI
+    app = gr.mount_gradio_app(app, demo, path="/gradio")
     
-    host = os.getenv("HOST", "0.0.0.0")
+    host = os.getenv("HOST", "localhost")
     port = int(os.getenv("PORT", "7860"))
     
     logger.info(f"Starting FastAPI server on {host}:{port}")
-    # logger.info(f"Gradio UI available at: http://{host}:{port}/gradio")  # Remove if Gradio is removed
+    logger.info(f"Gradio UI available at: http://{host}:{port}/gradio")  # Remove if Gradio is removed
     
     uvicorn.run(app, host=host, port=port, log_level="info", access_log=True)