refactor

app/main.py

@@ -1,93 +1,74 @@
-#CHATFED_ORCHESTRATOR
 import gradio as gr
-from fastapi import FastAPI, UploadFile, File, Form, Request
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.responses import StreamingResponse
 from langserve import add_routes
 from langgraph.graph import StateGraph, START, END
-from typing import Optional
 import uvicorn
 import os
 from datetime import datetime
 import logging
-from contextlib import asynccontextmanager
 from langchain_core.runnables import RunnableLambda
 import asyncio
 import json
-from functools import partial
 import base64
 
 from utils import getconfig
-from nodes import detect_file_type_node, ingest_node, geojson_direct_result_node, retrieve_node, generate_node_streaming, route_workflow, process_query_streaming
+from nodes import (
+    detect_file_type_node, ingest_node, geojson_direct_result_node, 
+    retrieve_node, generate_node_streaming, route_workflow, process_query_streaming
+)
 from models import GraphState, ChatUIInput, ChatUIFileInput
 
 config = getconfig("params.cfg")
-RETRIEVER = config.get("retriever", "RETRIEVER", fallback="https://giz-chatfed-retriever.hf.space")
-GENERATOR = config.get("generator", "GENERATOR", fallback="https://giz-chatfed-generator.hf.space")
-INGESTOR = config.get("ingestor", "INGESTOR", fallback="https://mtyrrell-chatfed-ingestor.hf.space")
-GEOJSON_INGESTOR = config.get("ingestor", "GEOJSON_INGESTOR", fallback="https://giz-eudr-chatfed-ingestor.hf.space")
-MAX_CONTEXT_CHARS = config.get("general", "MAX_CONTEXT_CHARS")
-
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 
-
 #----------------------------------------
-# CORE WORKFLOW GRAPH
+# LANGGRAPH WORKFLOW SETUP
 #----------------------------------------
-# graph setup
+# NOTE: Currently using manual node chaining for stability.
+# This graph is prepared for future agentic workflow expansion.
+
 workflow = StateGraph(GraphState)
 workflow.add_node("detect_file_type", detect_file_type_node)
 workflow.add_node("ingest", ingest_node)
 workflow.add_node("geojson_direct", geojson_direct_result_node)
 workflow.add_node("retrieve", retrieve_node)
-workflow.add_node("generate", generate_node_streaming) # Changed to generate_node_streaming
+workflow.add_node("generate", generate_node_streaming)
 
-# Add edges
 workflow.add_edge(START, "detect_file_type")
 workflow.add_edge("detect_file_type", "ingest")
 
-# Conditional routing after ingestion
 workflow.add_conditional_edges(
     "ingest",
     route_workflow,
-    {
-        "geojson_direct": "geojson_direct",
-        "standard": "retrieve"
-    }
+    {"geojson_direct": "geojson_direct", "standard": "retrieve"}
 )
 
-# Standard workflow
 workflow.add_edge("retrieve", "generate")
 workflow.add_edge("generate", END)
-
-# GeoJSON direct workflow
 workflow.add_edge("geojson_direct", END)
 
 compiled_graph = workflow.compile()
+# Future: Replace manual node chaining with: compiled_graph.astream(initial_state)
 
 
+#----------------------------------------
+# CHATUI ADAPTERS
+#----------------------------------------
+
 async def chatui_adapter(data):
-    """Updated to return content without SSE formatting for LangServe"""
+    """Text-only adapter for ChatUI"""
     try:
-        # Handle both dict and Pydantic model input
-        if hasattr(data, 'text'):
-            text = data.text
-        elif isinstance(data, dict) and 'text' in data:
-            text = data['text']
-        else:
-            logger.error(f"Unexpected input structure: {data}")
-            yield "Error: Invalid input format"
-            return
+        text = data.text if hasattr(data, 'text') else data.get('text', '')
         
-        # Collect all content and sources
         full_response = ""
         sources_collected = None
         
-        # Use the streaming function and return content directly (no SSE formatting)
         async for result in process_query_streaming(
             query=text,
-            file_upload=None,  # No file upload for text-only ChatUI
+            file_upload=None,
             reports_filter="",
             sources_filter="",
             subtype_filter="",
@@ -98,26 +79,19 @@ async def chatui_adapter(data):
                 content = result.get("content", "")
                 
                 if result_type == "data":
-                    # Accumulate the response text
                     full_response += content
                     yield content
-                    
                 elif result_type == "sources":
-                    # Store sources for later
                     sources_collected = content
-                    
                 elif result_type == "end":
-                    # Append sources to the final response if we have them
                     if sources_collected:
                         sources_text = "\n\n**Sources:**\n"
                         for i, source in enumerate(sources_collected, 1):
                             sources_text += f"{i}. [{source.get('title', 'Unknown')}]({source.get('link', '#')})\n"
                         yield sources_text
-                    
                 elif result_type == "error":
                     yield f"Error: {content}"
             else:
-                # Fallback for plain string results
                 yield str(result)
             
             await asyncio.sleep(0)
@@ -127,146 +101,26 @@ async def chatui_adapter(data):
         yield f"Error: {str(e)}"
 
 
-async def process_query_streaming_with_file_content(
-    query: str, 
-    file_content: Optional[bytes] = None,
-    filename: Optional[str] = None,
-    reports_filter: str = "", 
-    sources_filter: str = "", 
-    subtype_filter: str = "", 
-    year_filter: str = "", 
-    output_format: str = "structured"
-):
-    """
-    Modified streaming function that accepts file_content directly instead of file_upload
-    """
-    start_time = datetime.now()
-    session_id = f"chatui_{start_time.strftime('%Y%m%d_%H%M%S')}"
-    
-    try:
-        # Process ingestion first (non-streaming)
-        initial_state = {
-            "query": query,
-            "context": "",
-            "ingestor_context": "",
-            "result": "",
-            "sources": [],
-            "reports_filter": reports_filter or "",
-            "sources_filter": sources_filter or "",
-            "subtype_filter": subtype_filter or "",
-            "year_filter": year_filter or "",
-            "file_content": file_content,
-            "filename": filename,
-            "file_type": "unknown",
-            "workflow_type": "standard",
-            "metadata": {
-                "session_id": session_id,
-                "start_time": start_time.isoformat(),
-                "has_file_attachment": file_content is not None
-            }
-        }
-        
-        # Detect file type - merge the returned state with initial state
-        state_after_detect = {**initial_state, **detect_file_type_node(initial_state)}
-        
-        # Ingest if file provided - merge the returned state
-        state_after_ingest = {**state_after_detect, **ingest_node(state_after_detect)}
-        
-        # Route workflow
-        workflow_type = route_workflow(state_after_ingest)
-        
-        if workflow_type == "geojson_direct":
-            # For GeoJSON, return direct result
-            final_state = geojson_direct_result_node(state_after_ingest)
-            if output_format == "structured":
-                yield {"type": "data", "content": final_state["result"]}
-                yield {"type": "end", "content": ""}
-            else:
-                yield final_state["result"]
-        else:
-            # For standard workflow, retrieve first - merge the returned state
-            state_after_retrieve = {**state_after_ingest, **retrieve_node(state_after_ingest)}
-            
-            # Initialize variables for both output formats
-            sources_collected = None
-            accumulated_response = "" if output_format == "gradio" else None
-            
-            # Then stream generation
-            async for partial_state in generate_node_streaming(state_after_retrieve):
-                if "result" in partial_state:
-                    if output_format == "structured":
-                        yield {"type": "data", "content": partial_state["result"]}
-                    else:
-                        # Accumulate the content and yield the full accumulated response
-                        accumulated_response += partial_state["result"]
-                        yield accumulated_response
-                
-                # Collect sources for later
-                if "sources" in partial_state:
-                    sources_collected = partial_state["sources"]
-            
-            # Handle sources based on output format
-            if sources_collected:
-                if output_format == "structured":
-                    yield {"type": "sources", "content": sources_collected}
-                else:
-                    # Append sources to accumulated response
-                    sources_text = "\n\n**Sources:**\n"
-                    for i, source in enumerate(sources_collected, 1):
-                        if isinstance(source, dict):
-                            title = source.get('title', 'Unknown')
-                            link = source.get('link', '#')
-                            sources_text += f"{i}. [{title}]({link})\n"
-                        else:
-                            sources_text += f"{i}. {source}\n"
-                    
-                    accumulated_response += sources_text
-                    yield accumulated_response
-            
-            if output_format == "structured":
-                yield {"type": "end", "content": ""}
-        
-    except Exception as e:
-        logger.error(f"Streaming pipeline failed: {str(e)}")
-        if output_format == "structured":
-            yield {"type": "error", "content": f"Error: {str(e)}"}
-        else:
-            yield f"Error: {str(e)}"
-
-
 async def chatui_file_adapter(data):
-    """New adapter for file uploads with streaming response"""
+    """File upload adapter for ChatUI"""
     try:
-        logger.info(f"=== CHATUI FILE ADAPTER CALLED ===")
-        logger.info(f"Input data type: {type(data)}")
+        logger.info("=== CHATUI FILE ADAPTER CALLED ===")
         
-        # Handle both dict and Pydantic model input
-        if hasattr(data, 'text'):
-            text = data.text
-            files = getattr(data, 'files', None)
-        elif isinstance(data, dict):
-            text = data.get('text', '')
-            files = data.get('files', None)
-        else:
-            logger.error(f"Unexpected input structure: {data}")
-            yield "Error: Invalid input format"
-            return
+        text = data.text if hasattr(data, 'text') else data.get('text', '')
+        files = getattr(data, 'files', None) if hasattr(data, 'files') else data.get('files', None)
         
         logger.info(f"Text: {text[:100]}...")
         logger.info(f"Files present: {files is not None and len(files) > 0 if files else False}")
         
-        # Process file if provided
         file_content = None
         filename = None
         
         if files and len(files) > 0:
-            # Get the first file
             file_info = files[0]
             logger.info(f"Processing file: {file_info.get('name', 'unknown')}")
             
             if file_info.get('type') == 'base64' and file_info.get('content'):
                 try:
-                    # Decode base64 content
                     file_content = base64.b64decode(file_info['content'])
                     filename = file_info.get('name', 'uploaded_file')
                     logger.info(f"Decoded file: {filename}, size: {len(file_content)} bytes")
@@ -275,8 +129,8 @@ async def chatui_file_adapter(data):
                     yield f"Error: Failed to decode uploaded file - {str(e)}"
                     return
         
-        # Use the modified streaming function that handles file content directly
-        async for result in process_query_streaming_with_file_content(
+        # Use the unified streaming function
+        async for result in process_query_streaming(
             query=text,
             file_content=file_content,
             filename=filename,
@@ -284,7 +138,7 @@ async def chatui_file_adapter(data):
             sources_filter="",
             subtype_filter="",
             year_filter="",
-            output_format="structured"  # Use structured format for better control
+            output_format="structured"
         ):
             if isinstance(result, dict):
                 result_type = result.get("type", "data")
@@ -292,9 +146,7 @@ async def chatui_file_adapter(data):
                 
                 if result_type == "data":
                     yield content
-                    
                 elif result_type == "sources":
-                    # Format sources nicely
                     if content:
                         sources_text = "\n\n**Sources:**\n"
                         for i, source in enumerate(content, 1):
@@ -305,11 +157,9 @@ async def chatui_file_adapter(data):
                             else:
                                 sources_text += f"{i}. {source}\n"
                         yield sources_text
-                    
                 elif result_type == "error":
                     yield f"Error: {content}"
             else:
-                # Fallback for plain string results
                 yield str(result)
             
             await asyncio.sleep(0)
@@ -319,52 +169,11 @@ async def chatui_file_adapter(data):
         yield f"Error: {str(e)}"
 
 
-# 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.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")
-            
-            with gr.Column():
-                output = gr.Textbox(label="Response", lines=15, show_copy_button=True)
-
-        # Use streaming function
-        submit_btn.click(
-            fn=partial(process_query_streaming, output_format="gradio"),
-            inputs=[query_input, file_input, reports_filter_input, sources_filter_input, 
-                   subtype_filter_input, year_filter_input],
-            outputs=output,
-            show_progress="minimal"
-        )
-    
-    return demo
+#----------------------------------------
+# FASTAPI SETUP
+#----------------------------------------
 
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    logger.info("ChatFed Orchestrator starting up...")
-    yield
-    logger.info("Orchestrator shutting down...")
-
-app = FastAPI(
-    title="ChatFed Orchestrator",
-    version="1.0.0",
-    lifespan=lifespan,
-    docs_url=None,
-    redoc_url=None
-)
+app = FastAPI(title="ChatFed Orchestrator", version="1.0.0")
 
 @app.get("/health")
 async def health_check():
@@ -376,144 +185,33 @@ async def root():
         "message": "ChatFed Orchestrator API",
         "endpoints": {
             "health": "/health",
-            "chatfed-ui-stream": "/chatfed-ui-stream",  
-            "chatfed-with-file": "/chatfed-with-file",
-            "chatfed-with-file-stream": "/chatfed-with-file-stream",  # New Langserve route
+            "chatfed-ui-stream": "/chatfed-ui-stream (LangServe)",  
+            "chatfed-with-file": "/chatfed-with-file (FastAPI/SSE)",
+            "chatfed-with-file-stream": "/chatfed-with-file-stream (LangServe)",
+            "gradio": "/gradio"
         }
     }
 
 
-
-# # FILE UPLOAD ADAPTER
-async def chatfed_with_file_adapter(
-    query: str,
-    file_content: Optional[bytes] = None,
-    filename: Optional[str] = None,
-    reports_filter: str = "",
-    sources_filter: str = "",
-    subtype_filter: str = "",
-    year_filter: str = "",
-    session_id: Optional[str] = None,
-    user_id: Optional[str] = None
-):
-    """Async streaming adapter for file uploads"""
-    try:
-        # Use the same streaming logic as the working text endpoint
-        start_time = datetime.now()
-        if not session_id:
-            current_session_id = f"api_{start_time.strftime('%Y%m%d_%H%M%S')}"
-        else:
-            current_session_id = session_id
-        
-        # Create initial state
-        initial_state = {
-            "query": query,
-            "context": "",
-            "ingestor_context": "",
-            "result": "",
-            "sources": [],
-            "reports_filter": reports_filter or "",
-            "sources_filter": sources_filter or "",
-            "subtype_filter": subtype_filter or "",
-            "year_filter": year_filter or "",
-            "file_content": file_content,
-            "filename": filename,
-            "file_type": "unknown",
-            "workflow_type": "standard",
-            "metadata": {
-                "session_id": current_session_id,
-                "user_id": user_id,
-                "start_time": start_time.isoformat(),
-                "has_file_attachment": file_content is not None
-            }
-        }
-        
-        # Process non-streaming steps first
-        state_after_detect = {**initial_state, **detect_file_type_node(initial_state)}
-        state_after_ingest = {**state_after_detect, **ingest_node(state_after_detect)}
-        
-        # Route workflow
-        workflow_type = route_workflow(state_after_ingest)
-        
-        if workflow_type == "geojson_direct":
-            # For GeoJSON, return direct result
-            final_state = geojson_direct_result_node(state_after_ingest)
-            yield final_state["result"]
-        else:
-            # For standard workflow, retrieve first
-            state_after_retrieve = {**state_after_ingest, **retrieve_node(state_after_ingest)}
-            
-            # Then stream generation
-            async for partial_state in generate_node_streaming(state_after_retrieve):
-                if "result" in partial_state:
-                    yield partial_state["result"]
-                await asyncio.sleep(0)  # Make it properly async
-        
-    except Exception as e:
-        logger.error(f"File upload streaming failed: {str(e)}")
-        yield f"Error: {str(e)}"
-
-# NON-STREAMING FILE UPLOAD
-# @app.post("/chatfed-with-file")
-# async def chatfed_with_file(
-#     query: str = Form(...),
-#     file: Optional[UploadFile] = File(None),
-#     reports_filter: Optional[str] = Form(""),
-#     sources_filter: Optional[str] = Form(""),
-#     subtype_filter: Optional[str] = Form(""),
-#     year_filter: Optional[str] = Form(""),
-#     session_id: Optional[str] = Form(None),
-#     user_id: Optional[str] = Form(None)
-# ):
-#     """Endpoint for queries with optional file attachments + streaming"""
-    
-#     # Read file content first
-#     file_content = None
-#     filename = None
-    
-#     if file:
-#         file_content = await file.read()
-#         filename = file.filename
-    
-#     # Stream the response instead of collecting chunks
-#     async def stream_generator():
-#         async for chunk in chatfed_with_file_adapter(
-#             query=query,
-#             file_content=file_content,
-#             filename=filename,
-#             reports_filter=reports_filter,
-#             sources_filter=sources_filter,
-#             subtype_filter=subtype_filter,
-#             year_filter=year_filter,
-#             session_id=session_id,
-#             user_id=user_id
-#         ):
-#             yield chunk
-    
-#     return StreamingResponse(
-#         stream_generator(),
-#         media_type="text/plain"
-#     )
-
-# MAIN FILE UPLOAD STREAMING ENDPOINT
 @app.post("/chatfed-with-file")
-async def chatfed_with_file_stream(
+async def chatfed_with_file_endpoint(
     query: str = Form(...),
-    file: Optional[UploadFile] = File(None),
-    reports_filter: Optional[str] = Form(""),
-    sources_filter: Optional[str] = Form(""),
-    subtype_filter: Optional[str] = Form(""),
-    year_filter: Optional[str] = Form(""),
-    session_id: Optional[str] = Form(None),
-    user_id: Optional[str] = Form(None)
+    file: UploadFile = File(None),
+    reports_filter: str = Form(""),
+    sources_filter: str = Form(""),
+    subtype_filter: str = Form(""),
+    year_filter: str = Form(""),
+    session_id: str = Form(None),
+    user_id: str = Form(None)
 ):
-    """File upload endpoint with proper SSE streaming for ChatUI"""
-    
-    logger.info(f"=== FILE UPLOAD ENDPOINT CALLED ===")
+    """
+    File upload endpoint with SSE streaming for ChatUI.
+    This endpoint is currently used by ChatUI's fileUploadUrl config.
+    """
+    logger.info("=== FILE UPLOAD ENDPOINT CALLED ===")
     logger.info(f"Query: {query[:100]}...")
     logger.info(f"File: {file.filename if file else 'None'}")
     
-    # Read file content
     file_content = None
     filename = None
     
@@ -522,11 +220,11 @@ async def chatfed_with_file_stream(
         filename = file.filename
     
     async def sse_generator():
-        """Generate Server-Sent Events format for ChatUI"""
+        """Generate Server-Sent Events for ChatUI"""
         try:
             token_id = 0
             
-            async for chunk in chatfed_with_file_adapter(
+            async for chunk in process_query_streaming(
                 query=query,
                 file_content=file_content,
                 filename=filename,
@@ -534,24 +232,34 @@ async def chatfed_with_file_stream(
                 sources_filter=sources_filter,
                 subtype_filter=subtype_filter,
                 year_filter=year_filter,
-                session_id=session_id,
-                user_id=user_id
+                output_format="structured"
             ):
-                if isinstance(chunk, str) and chunk.strip():
-                    # Format as SSE data that ChatUI expects
-                    token_data = {
-                        "token": chunk,
-                        "text": chunk,
-                        "content": chunk
-                    }
-                    yield f"data: {json.dumps(token_data)}\n\n"
-                    token_id += 1
+                if isinstance(chunk, dict):
+                    chunk_type = chunk.get("type", "data")
+                    content = chunk.get("content", "")
                     
+                    if chunk_type == "data" and content:
+                        token_data = {"token": content, "text": content, "content": content}
+                        yield f"data: {json.dumps(token_data)}\n\n"
+                        token_id += 1
+                    elif chunk_type == "sources" and content:
+                        # Format sources for display
+                        sources_text = "\n\n**Sources:**\n"
+                        for i, source in enumerate(content, 1):
+                            if isinstance(source, dict):
+                                title = source.get('title', 'Unknown')
+                                link = source.get('link', '#')
+                                sources_text += f"{i}. [{title}]({link})\n"
+                        token_data = {"token": sources_text, "text": sources_text, "content": sources_text}
+                        yield f"data: {json.dumps(token_data)}\n\n"
+                    elif chunk_type == "error":
+                        error_data = {"error": content}
+                        yield f"data: {json.dumps(error_data)}\n\n"
+                        
                 await asyncio.sleep(0)
             
-            # Send end marker
             yield f"data: [DONE]\n\n"
-            logger.info("Generator stream ended")
+            logger.info("SSE stream completed")
             
         except Exception as e:
             logger.error(f"SSE generation error: {str(e)}")
@@ -565,11 +273,15 @@ async def chatfed_with_file_stream(
             "Cache-Control": "no-cache",
             "Connection": "keep-alive",
             "Access-Control-Allow-Origin": "*",
-            "Access-Control-Allow-Headers": "*",
         }
     )
 
-# Add the existing text-only Langserve route
+
+#----------------------------------------
+# LANGSERVE ROUTES
+#----------------------------------------
+
+# Text-only endpoint
 add_routes(
     app,
     RunnableLambda(chatui_adapter),
@@ -580,7 +292,7 @@ add_routes(
     enable_public_trace_link_endpoint=True,
 )
 
-# Add the new file upload Langserve route
+# File upload endpoint (LangServe version for future migration)
 add_routes(
     app,
     RunnableLambda(chatui_file_adapter),
@@ -591,17 +303,72 @@ add_routes(
     enable_public_trace_link_endpoint=True,
 )
 
+
+#----------------------------------------
+# GRADIO INTERFACE
+#----------------------------------------
+
+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.")
+        
+        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 = gr.Textbox(label="Reports Filter", placeholder="e.g., annual_reports")
+                    sources_filter = gr.Textbox(label="Sources Filter", placeholder="e.g., internal")
+                    subtype_filter = gr.Textbox(label="Subtype Filter", placeholder="e.g., financial")
+                    year_filter = gr.Textbox(label="Year Filter", placeholder="e.g., 2024")
+                
+                submit_btn = gr.Button("Submit", variant="primary")
+            
+            with gr.Column():
+                output = gr.Textbox(label="Response", lines=15, show_copy_button=True)
+
+        async def gradio_handler(query, file, reports, sources, subtype, year):
+            """Handler for Gradio interface"""
+            result = ""
+            async for chunk in process_query_streaming(
+                query=query,
+                file_upload=file,
+                reports_filter=reports,
+                sources_filter=sources,
+                subtype_filter=subtype,
+                year_filter=year,
+                output_format="gradio"
+            ):
+                result = chunk  # Each chunk is the full accumulated text
+                yield result
+        
+        submit_btn.click(
+            fn=gradio_handler,
+            inputs=[query_input, file_input, reports_filter, sources_filter, subtype_filter, year_filter],
+            outputs=output,
+        )
+    
+    return demo
+
+
+#----------------------------------------
+# MAIN
+#----------------------------------------
+
 if __name__ == "__main__":
-    # Create Gradio interface
     demo = create_gradio_interface()
-    
-    # Mount Gradio app to FastAPI
     app = gr.mount_gradio_app(app, demo, path="/gradio")
     
     host = os.getenv("HOST", "0.0.0.0")  
     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") 
+    logger.info(f"Starting ChatFed Orchestrator on {host}:{port}")
+    logger.info(f"Gradio UI: http://{host}:{port}/gradio")
+    logger.info(f"API Docs: http://{host}:{port}/docs")
     
     uvicorn.run(app, host=host, port=port, log_level="info", access_log=True)

app/models.py

@@ -1,14 +1,14 @@
-# Models
 from typing import Optional, Dict, Any, List
 from typing_extensions import TypedDict
 from pydantic import BaseModel
 
 class GraphState(TypedDict):
+    """State object passed through LangGraph workflow"""
     query: str
     context: str
     ingestor_context: str
     result: str
-    sources: Optional[List[Dict[str, str]]]  # Added for ChatUI sources
+    sources: List[Dict[str, str]]  # Always present, no Optional needed
     reports_filter: str
     sources_filter: str
     subtype_filter: str
@@ -19,26 +19,12 @@ class GraphState(TypedDict):
     file_type: Optional[str]
     workflow_type: Optional[str]  # 'standard' or 'geojson_direct'
 
-class ChatFedInput(TypedDict):
-    query: str
-    reports_filter: Optional[str]
-    sources_filter: Optional[str] 
-    subtype_filter: Optional[str]
-    year_filter: Optional[str]
-    session_id: Optional[str]
-    user_id: Optional[str]
-    file_content: Optional[bytes]
-    filename: Optional[str]
-
-class ChatFedOutput(TypedDict):
-    result: str
-    metadata: Dict[str, Any]
-
 class ChatUIInput(BaseModel):
+    """Input model for text-only ChatUI requests"""
     text: str
 
-# New model for file upload support
 class ChatUIFileInput(BaseModel):
+    """Input model for ChatUI requests with file attachments"""
     text: str
     files: Optional[List[Dict[str, Any]]] = None
 

app/nodes.py

@@ -1,36 +1,32 @@
-from utils import detect_file_type, convert_context_to_list
+from utils import detect_file_type, convert_context_to_list, merge_state, getconfig
 from models import GraphState
 from datetime import datetime
 import tempfile
 import os
 from gradio_client import Client, file
 import logging
-from utils import getconfig
 import dotenv
-from typing_extensions import TypedDict
 import httpx
 import json
-from typing import Generator
+from typing import Generator, Optional
 
 dotenv.load_dotenv()
 
 logger = logging.getLogger(__name__)
 
+# Load config once at module level
 config = getconfig("params.cfg")
 RETRIEVER = config.get("retriever", "RETRIEVER", fallback="https://giz-chatfed-retriever.hf.space")
 GENERATOR = config.get("generator", "GENERATOR", fallback="https://giz-chatfed-generator.hf.space")
 INGESTOR = config.get("ingestor", "INGESTOR", fallback="https://mtyrrell-chatfed-ingestor.hf.space")
 GEOJSON_INGESTOR = config.get("ingestor", "GEOJSON_INGESTOR", fallback="https://giz-eudr-chatfed-ingestor.hf.space")
-MAX_CONTEXT_CHARS = config.get("general", "MAX_CONTEXT_CHARS")
+MAX_CONTEXT_CHARS = int(config.get("general", "MAX_CONTEXT_CHARS"))
 
-ingestor_url = INGESTOR
-retriever_url = RETRIEVER
-generator_url = GENERATOR
-geojson_ingestor_url = GEOJSON_INGESTOR
 
-# CORE PROCESSING NODES
 #----------------------------------------
-# File type detection node
+# LANGGRAPH NODE FUNCTIONS
+#----------------------------------------
+
 def detect_file_type_node(state: GraphState) -> GraphState:
     """Detect file type and determine workflow"""
     file_type = "unknown"
@@ -38,12 +34,7 @@ def detect_file_type_node(state: GraphState) -> GraphState:
     
     if state.get("file_content") and state.get("filename"):
         file_type = detect_file_type(state["filename"], state["file_content"])
-        
-        # Determine workflow based on file type
-        if file_type == "geojson":
-            workflow_type = "geojson_direct"
-        else:
-            workflow_type = "standard"
+        workflow_type = "geojson_direct" if file_type == "geojson" else "standard"
     
     metadata = state.get("metadata", {})
     metadata.update({
@@ -57,12 +48,11 @@ def detect_file_type_node(state: GraphState) -> GraphState:
         "metadata": metadata
     }
 
-# Module functions
+
 def ingest_node(state: GraphState) -> GraphState:
     """Process file through appropriate ingestor based on file type"""
     start_time = datetime.now()
     
-    # If no file provided, skip this step
     if not state.get("file_content") or not state.get("filename"):
         logger.info("No file provided, skipping ingestion")
         return {"ingestor_context": "", "metadata": state.get("metadata", {})}
@@ -72,35 +62,23 @@ def ingest_node(state: GraphState) -> GraphState:
     
     try:
         # Choose ingestor based on file type
-        if file_type == "geojson":
-            ingestor_url = GEOJSON_INGESTOR
-            logger.info(f"Using GeoJSON ingestor: {ingestor_url}")
-        else:
-            ingestor_url = INGESTOR
-            logger.info(f"Using standard ingestor: {ingestor_url}")
+        ingestor_url = GEOJSON_INGESTOR if file_type == "geojson" else INGESTOR
+        logger.info(f"Using ingestor: {ingestor_url}")
         
         client = Client(ingestor_url, hf_token=os.getenv("HF_TOKEN"))
         
-        # Create a temporary file to upload
+        # Create temporary file for upload
         with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(state["filename"])[1]) as tmp_file:
             tmp_file.write(state["file_content"])
             tmp_file_path = tmp_file.name
         
         try:
-            # Call the ingestor's ingest endpoint
-            ingestor_context = client.predict(
-                file(tmp_file_path),
-                api_name="/ingest"
-            )
-            
+            ingestor_context = client.predict(file(tmp_file_path), api_name="/ingest")
             logger.info(f"Ingest result length: {len(ingestor_context) if ingestor_context else 0}")
             
-            # Handle error cases
             if isinstance(ingestor_context, str) and ingestor_context.startswith("Error:"):
                 raise Exception(ingestor_context)
-            
         finally:
-            # Clean up temporary file
             os.unlink(tmp_file_path)
         
         duration = (datetime.now() - start_time).total_seconds()
@@ -112,10 +90,7 @@ def ingest_node(state: GraphState) -> GraphState:
             "ingestor_used": ingestor_url
         })
         
-        return {
-            "ingestor_context": ingestor_context,
-            "metadata": metadata
-        }
+        return {"ingestor_context": ingestor_context, "metadata": metadata}
         
     except Exception as e:
         duration = (datetime.now() - start_time).total_seconds()
@@ -129,13 +104,12 @@ def ingest_node(state: GraphState) -> GraphState:
         })
         return {"ingestor_context": "", "metadata": metadata}
 
+
 def geojson_direct_result_node(state: GraphState) -> GraphState:
-    """For GeoJSON files, return ingestor results directly without retrieval/generation"""
+    """For GeoJSON files, return ingestor results directly"""
     logger.info("Processing GeoJSON file - returning direct results")
     
     ingestor_context = state.get("ingestor_context", "")
-    
-    # For GeoJSON files, the ingestor result is the final result
     result = ingestor_context if ingestor_context else "No results from GeoJSON processing."
     
     metadata = state.get("metadata", {})
@@ -144,12 +118,11 @@ def geojson_direct_result_node(state: GraphState) -> GraphState:
         "result_length": len(result)
     })
     
-    return {
-        "result": result,
-        "metadata": metadata
-    }
+    return {"result": result, "metadata": metadata}
+
 
 def retrieve_node(state: GraphState) -> GraphState:
+    """Retrieve relevant context from vector store"""
     start_time = datetime.now()
     logger.info(f"Retrieval: {state['query'][:50]}...")
     
@@ -187,38 +160,29 @@ def retrieve_node(state: GraphState) -> GraphState:
         return {"context": "", "metadata": metadata}
 
 
-
-
-# MAIN STREAMING GENERATOR
 async def generate_node_streaming(state: GraphState) -> Generator[GraphState, None, None]:
-    """Streaming version that calls generator's FastAPI endpoint"""
+    """Streaming generation using generator's FastAPI endpoint"""
     start_time = datetime.now()
     logger.info(f"Generation (streaming): {state['query'][:50]}...")
     
     try:
-        # Get MAX_CONTEXT_CHARS at the beginning so it's available throughout the function
-        MAX_CONTEXT_CHARS = int(config.get("general", "MAX_CONTEXT_CHARS"))
-        logger.info(f"Using MAX_CONTEXT_CHARS: {MAX_CONTEXT_CHARS}")
-        
-        # Combine retriever context with ingestor context
+        # Combine contexts
         retrieved_context = state.get("context", "")
         ingestor_context = state.get("ingestor_context", "")
         
-        logger.info(f"Original context lengths - Ingestor: {len(ingestor_context)}, Retrieved: {len(retrieved_context)}")
+        logger.info(f"Context lengths - Ingestor: {len(ingestor_context)}, Retrieved: {len(retrieved_context)}")
         
-        # Convert contexts to list format expected by generator
+        # Build context list with truncation
         context_list = []
         total_context_chars = 0
         
         if ingestor_context:
-            # Truncate ingestor context if it's too long
-            if len(ingestor_context) > MAX_CONTEXT_CHARS:
-                logger.warning(f"Truncating ingestor context from {len(ingestor_context)} to {MAX_CONTEXT_CHARS} characters")
-                truncated_ingestor = ingestor_context[:MAX_CONTEXT_CHARS] + "...\n[Content truncated due to length]"
-            else:
-                truncated_ingestor = ingestor_context
+            truncated_ingestor = (
+                ingestor_context[:MAX_CONTEXT_CHARS] + "...\n[Content truncated due to length]"
+                if len(ingestor_context) > MAX_CONTEXT_CHARS
+                else ingestor_context
+            )
             
-            # Add ingestor context
             context_list.append({
                 "answer": truncated_ingestor,
                 "answer_metadata": {
@@ -231,154 +195,101 @@ async def generate_node_streaming(state: GraphState) -> Generator[GraphState, No
             total_context_chars += len(truncated_ingestor)
         
         if retrieved_context and total_context_chars < MAX_CONTEXT_CHARS:
-            # Convert retrieved context to list and add
             retrieved_list = convert_context_to_list(retrieved_context)
-            
-            # Add retrieved context items until we hit the limit
             remaining_chars = MAX_CONTEXT_CHARS - total_context_chars
+            
             for item in retrieved_list:
                 item_text = item.get("answer", "")
                 if len(item_text) <= remaining_chars:
                     context_list.append(item)
                     remaining_chars -= len(item_text)
                 else:
-                    # Truncate this item and stop
-                    if remaining_chars > 100:  # Only add if we have meaningful space left
+                    if remaining_chars > 100:
                         item["answer"] = item_text[:remaining_chars-50] + "...\n[Content truncated]"
                         context_list.append(item)
                     break
         
-        # Calculate final context size
         final_context_size = sum(len(item.get("answer", "")) for item in context_list)
         logger.info(f"Final context size: {final_context_size} characters (limit: {MAX_CONTEXT_CHARS})")
         
-        # Prepare the request payload
-        payload = {
-            "query": state["query"],
-            "context": context_list
-        }
+        payload = {"query": state["query"], "context": context_list}
         
-        # Determine generator URL - handle both Hugging Face and direct URLs
+        # Normalize generator URL
         generator_url = GENERATOR
-        
         if not generator_url.startswith('http'):
-            # Allows for easy specification of space in config (converts to URL)
-            # Replace '/' with '-' for Hugging Face space URLs
-            # Force the replacement to ensure it works
-            space_name = generator_url.replace('/', '-').replace('_', '-')
+            space_name = generator_url.replace('/', '-')
             generator_url = f"https://{space_name}.hf.space"
-    
-        
-        # Try FastAPI endpoint first, fallback to Gradio if needed
-        fastapi_success = False
         
-        try:
-            # Make streaming request to generator's FastAPI endpoint
-            async with httpx.AsyncClient(timeout=300.0, verify=False) as client:
+        # Stream from generator
+        async with httpx.AsyncClient(timeout=300.0, verify=False) as client:
+            async with client.stream(
+                "POST",
+                f"{generator_url}/generate/stream",
+                json=payload,
+                headers={"Content-Type": "application/json"}
+            ) as response:
+                if response.status_code != 200:
+                    raise Exception(f"Generator returned status {response.status_code}")
                 
-                async with client.stream(
-                    "POST",
-                    f"{generator_url}/generate/stream",
-                    json=payload,
-                    headers={"Content-Type": "application/json"}
-                ) as response:
-                    if response.status_code != 200:
-                        error_text = await response.aread()
-                        raise Exception(f"FastAPI endpoint returned status {response.status_code}")
-                    
-                    current_text = ""
-                    sources = None
-                    event_type = None
+                current_text = ""
+                sources = None
+                event_type = None
+                
+                async for line in response.aiter_lines():
+                    if not line.strip():
+                        continue
                     
-                    async for line in response.aiter_lines():
-                        if not line.strip():
-                            continue
+                    if line.startswith("event: "):
+                        event_type = line[7:].strip()
+                        continue
+                    elif line.startswith("data: "):
+                        data_content = line[6:].strip()
                         
-                        # Parse SSE format
-                        if line.startswith("event: "):
-                            event_type = line[7:].strip()
-                            continue
-                        elif line.startswith("data: "):
-                            data_content = line[6:].strip()
-                            
-                            if event_type == "data":
-                                # Text chunk
-                                try:
-                                    chunk = json.loads(data_content)
-                                    if isinstance(chunk, str):
-                                        current_text += chunk
-                                        
-                                        metadata = state.get("metadata", {})
-                                        metadata.update({
-                                            "generation_duration": (datetime.now() - start_time).total_seconds(),
-                                            "result_length": len(current_text),
-                                            "generation_success": True,
-                                            "streaming": True,
-                                            "generator_type": "fastapi",
-                                            "context_chars_used": final_context_size
-                                        })
-                                        
-                                        yield {
-                                            "result": chunk,  # Send only the new chunk
-                                            "metadata": metadata
-                                        }
-                                except json.JSONDecodeError:
-                                    # Handle plain text chunks
-                                    current_text += data_content
-                                    
-                                    metadata = state.get("metadata", {})
-                                    metadata.update({
-                                        "generation_duration": (datetime.now() - start_time).total_seconds(),
-                                        "result_length": len(current_text),
-                                        "generation_success": True,
-                                        "streaming": True,
-                                        "generator_type": "fastapi",
-                                        "context_chars_used": final_context_size
-                                    })
-                                    
-                                    yield {
-                                        "result": data_content,
-                                        "metadata": metadata
-                                    }
-                            
-                            elif event_type == "sources":
-                                # Sources data
-                                try:
-                                    sources_data = json.loads(data_content)
-                                    sources = sources_data.get("sources", [])
-                                    
-                                    # Update state with sources
-                                    metadata = state.get("metadata", {})
-                                    metadata.update({
-                                        "sources_received": True,
-                                        "sources_count": len(sources)
-                                    })
-                                    
-                                    yield {
-                                        "sources": sources,
-                                        "metadata": metadata
-                                    }
-                                except json.JSONDecodeError:
-                                    logger.warning(f"Failed to parse sources data: {data_content}")
+                        if event_type == "data":
+                            try:
+                                chunk = json.loads(data_content)
+                                if isinstance(chunk, str):
+                                    current_text += chunk
+                            except json.JSONDecodeError:
+                                current_text += data_content
+                                chunk = data_content
                             
-                            elif event_type == "end":
-                                # Stream ended
-                                logger.info("Generator stream ended")
-                                fastapi_success = True
-                                break
+                            metadata = state.get("metadata", {})
+                            metadata.update({
+                                "generation_duration": (datetime.now() - start_time).total_seconds(),
+                                "result_length": len(current_text),
+                                "generation_success": True,
+                                "streaming": True,
+                                "context_chars_used": final_context_size
+                            })
                             
-                            elif event_type == "error":
-                                # Error occurred
-                                try:
-                                    error_data = json.loads(data_content)
-                                    raise Exception(error_data.get("error", "Unknown error"))
-                                except json.JSONDecodeError:
-                                    raise Exception(data_content)
-        
-        except Exception as fastapi_error:
-            # Handle FastAPI-specific errors 
-            logger.warning(f"FastAPI endpoint failed: {str(fastapi_error)}")
-            raise fastapi_error
+                            yield {"result": chunk, "metadata": metadata}
+                        
+                        elif event_type == "sources":
+                            try:
+                                sources_data = json.loads(data_content)
+                                sources = sources_data.get("sources", [])
+                                
+                                metadata = state.get("metadata", {})
+                                metadata.update({
+                                    "sources_received": True,
+                                    "sources_count": len(sources)
+                                })
+                                
+                                yield {"sources": sources, "metadata": metadata}
+                            except json.JSONDecodeError:
+                                logger.warning(f"Failed to parse sources: {data_content}")
+                        
+                        elif event_type == "end":
+                            logger.info("Generator stream ended")
+                            break
+                        
+                        elif event_type == "error":
+                            try:
+                                error_data = json.loads(data_content)
+                                raise Exception(error_data.get("error", "Unknown error"))
+                            except json.JSONDecodeError:
+                                raise Exception(data_content)
         
     except Exception as e:
         duration = (datetime.now() - start_time).total_seconds()
@@ -393,27 +304,38 @@ async def generate_node_streaming(state: GraphState) -> Generator[GraphState, No
         })
         yield {"result": f"Error: {str(e)}", "metadata": metadata}
 
-# Conditional routing function
-def route_workflow(state: GraphState) -> str:
-    """Route to appropriate workflow based on file type"""
-    workflow_type = state.get("workflow_type", "standard")
-    return workflow_type
 
+def route_workflow(state: GraphState) -> str:
+    """Conditional routing based on workflow type"""
+    return state.get("workflow_type", "standard")
 
 
+#----------------------------------------
+# UNIFIED STREAMING PROCESSOR
+#----------------------------------------
 
-async def process_query_streaming(query: str, file_upload, reports_filter: str = "", sources_filter: str = "", 
-                                 subtype_filter: str = "", year_filter: str = "", 
-                                 output_format: str = "structured"):
+async def process_query_streaming(
+    query: str, 
+    file_upload=None,
+    file_content: Optional[bytes] = None,
+    filename: Optional[str] = None,
+    reports_filter: str = "", 
+    sources_filter: str = "", 
+    subtype_filter: str = "", 
+    year_filter: str = "", 
+    output_format: str = "structured"
+):
     """
-    Unified streaming function that yields partial results
+    Unified streaming function supporting both file objects and raw content.
     
     Args:
-        output_format: "structured" for dict format, "gradio" for plain text format
+        query: User query string
+        file_upload: File object from Gradio (optional)
+        file_content: Raw file bytes (optional, alternative to file_upload)
+        filename: Filename for raw content (required if file_content provided)
+        output_format: "structured" returns dicts, "gradio" returns accumulated text
     """
-    file_content = None
-    filename = None
-    
+    # Handle file_upload if provided (Gradio use case)
     if file_upload is not None:
         try:
             with open(file_upload.name, 'rb') as f:
@@ -429,10 +351,10 @@ async def process_query_streaming(query: str, file_upload, reports_filter: str =
             return
     
     start_time = datetime.now()
-    session_id = f"gradio_{start_time.strftime('%Y%m%d_%H%M%S')}"
+    session_id = f"stream_{start_time.strftime('%Y%m%d_%H%M%S')}"
     
     try:
-        # Process ingestion first (non-streaming)
+        # Build initial state
         initial_state = {
             "query": query,
             "context": "",
@@ -454,51 +376,41 @@ async def process_query_streaming(query: str, file_upload, reports_filter: str =
             }
         }
         
-        # Detect file type - merge the returned state with initial state
-        state_after_detect = {**initial_state, **detect_file_type_node(initial_state)}
-        
-        # Ingest if file provided - merge the returned state
-        state_after_ingest = {**state_after_detect, **ingest_node(state_after_detect)}
+        # Execute workflow nodes
+        state = merge_state(initial_state, detect_file_type_node(initial_state))
+        state = merge_state(state, ingest_node(state))
         
-        # Route workflow
-        workflow_type = route_workflow(state_after_ingest)
+        workflow_type = route_workflow(state)
         
         if workflow_type == "geojson_direct":
-            # For GeoJSON, return direct result
-            final_state = geojson_direct_result_node(state_after_ingest)
+            final_state = geojson_direct_result_node(state)
             if output_format == "structured":
                 yield {"type": "data", "content": final_state["result"]}
                 yield {"type": "end", "content": ""}
             else:
                 yield final_state["result"]
         else:
-            # For standard workflow, retrieve first - merge the returned state
-            state_after_retrieve = {**state_after_ingest, **retrieve_node(state_after_ingest)}
+            state = merge_state(state, retrieve_node(state))
             
-            # Initialize variables for both output formats
             sources_collected = None
             accumulated_response = "" if output_format == "gradio" else None
             
-            # Then stream generation
-            async for partial_state in generate_node_streaming(state_after_retrieve):
+            async for partial_state in generate_node_streaming(state):
                 if "result" in partial_state:
                     if output_format == "structured":
                         yield {"type": "data", "content": partial_state["result"]}
                     else:
-                        # Accumulate the content and yield the full accumulated response
                         accumulated_response += partial_state["result"]
                         yield accumulated_response
                 
-                # Collect sources for later
                 if "sources" in partial_state:
                     sources_collected = partial_state["sources"]
             
-            # Handle sources based on output format
+            # Format and yield sources
             if sources_collected:
                 if output_format == "structured":
                     yield {"type": "sources", "content": sources_collected}
                 else:
-                    # Append sources to accumulated response
                     sources_text = "\n\n**Sources:**\n"
                     for i, source in enumerate(sources_collected, 1):
                         if isinstance(source, dict):

app/utils.py

@@ -2,28 +2,24 @@ import configparser
 import logging
 import os
 import ast
-import re
+import json
 from dotenv import load_dotenv
 from typing import Optional, Dict, Any, List
+from models import GraphState
 
-# Local .env file
 load_dotenv()
 
 logger = logging.getLogger(__name__)
 
 def getconfig(configfile_path: str):
-    """
-    Read the config file
-    Params
-    ----------------
-    configfile_path: file path of .cfg file
-    """
+    """Read the config file"""
     config = configparser.ConfigParser()
     try:
         config.read_file(open(configfile_path))
         return config
     except:
         logging.warning("config file not found")
+        return None
 
 
 def get_auth(provider: str) -> dict:
@@ -33,7 +29,7 @@ def get_auth(provider: str) -> dict:
         "qdrant": {"api_key": os.getenv("QDRANT_API_KEY")},
     }
     
-    provider = provider.lower()  # Normalize to lowercase
+    provider = provider.lower()
     
     if provider not in auth_configs:
         raise ValueError(f"Unsupported provider: {provider}")
@@ -42,24 +38,22 @@ def get_auth(provider: str) -> dict:
     api_key = auth_config.get("api_key")
     
     if not api_key:
-        logging.warning(f"No API key found for provider '{provider}'. Please set the appropriate environment variable.")
+        logging.warning(f"No API key found for provider '{provider}'")
         auth_config["api_key"] = None
     
     return auth_config
 
-# File type detection
+
 def detect_file_type(filename: str, file_content: bytes = None) -> str:
     """Detect file type based on extension and content"""
     if not filename:
         return "unknown"
     
-    # Get file extension
     _, ext = os.path.splitext(filename.lower())
     
-    # Define file type mappings
     file_type_mappings = {
         '.geojson': 'geojson',
-        '.json': 'json',  # Could be geojson, will check content
+        '.json': 'json',
         '.pdf': 'text',
         '.docx': 'text',
         '.doc': 'text',
@@ -75,29 +69,28 @@ def detect_file_type(filename: str, file_content: bytes = None) -> str:
     # For JSON files, check if it's actually GeoJSON
     if detected_type == 'json' and file_content:
         try:
-            import json
             content_str = file_content.decode('utf-8')
             data = json.loads(content_str)
-            # Check if it has GeoJSON structure
-            if isinstance(data, dict) and ('type' in data and data.get('type') == 'FeatureCollection'):
+            if isinstance(data, dict) and data.get('type') == 'FeatureCollection':
                 detected_type = 'geojson'
-            elif isinstance(data, dict) and ('type' in data and data.get('type') in ['Feature', 'Point', 'LineString', 'Polygon', 'MultiPoint', 'MultiLineString', 'MultiPolygon', 'GeometryCollection']):
+            elif isinstance(data, dict) and data.get('type') in [
+                'Feature', 'Point', 'LineString', 'Polygon', 
+                'MultiPoint', 'MultiLineString', 'MultiPolygon', 'GeometryCollection'
+            ]:
                 detected_type = 'geojson'
         except:
-            pass  # Keep as json if parsing fails
+            pass
     
     logger.info(f"Detected file type: {detected_type} for file: {filename}")
     return detected_type
 
-# Helper function to convert retrieval context to expected format
+
 def convert_context_to_list(context: str) -> List[Dict[str, Any]]:
     """Convert string context to list format expected by generator"""
     try:
-        # Try to parse as list first
         if context.startswith('['):
             return ast.literal_eval(context)
         else:
-            # If it's a string, wrap it in a simple format
             return [{
                 "answer": context,
                 "answer_metadata": {
@@ -108,7 +101,6 @@ def convert_context_to_list(context: str) -> List[Dict[str, Any]]:
                 }
             }]
     except:
-        # Fallback: simple string wrapping
         return [{
             "answer": context,
             "answer_metadata": {
@@ -117,4 +109,9 @@ def convert_context_to_list(context: str) -> List[Dict[str, Any]]:
                 "year": "Unknown",
                 "source": "Retriever"
             }
-        }]
+        }]
+
+
+def merge_state(base_state: GraphState, updates: dict) -> GraphState:
+    """Helper to merge node updates into base state"""
+    return {**base_state, **updates}