ea4all-gradio-agents-mcp-hackathon-tools-refactoring-vqa

ea4all/ea4all_mcp.py

@@ -13,6 +13,9 @@ import ea4all.src.ea4all_vqa.graph as e4v
 import ea4all.src.ea4all_gra.graph as e4t
 import ea4all.src.shared.utils as e4u
 from ea4all.src.graph import super_graph
+from ea4all.src.ea4all_apm.graph import apm_graph
+from ea4all.src.ea4all_vqa.graph import diagram_graph
+
 #from ea4all.src.pmo_crew.crew_runner import run_pmo_crew
 
 from typing import AsyncGenerator
@@ -66,100 +69,11 @@ async def run_qna_agentic_system(question: str) -> AsyncGenerator[list, None]:
     if not question:
         format_response = "Hi, how are you today? To start using the EA4ALL MCP Tool, provide the required Inputs!"
         chat_memory.append(ChatMessage(role="assistant", content=format_response))
-        yield chat_memory
-
-    if not chat_memory:
-        chat_memory.append(ChatMessage(role="user", content=question))
-        yield chat_memory
-
-    if question:
-        #capture user ip
-        #ea4all_user = e4u.get_user_identification(request)
-
-        ##Initialise APM Graph
-        #apm_graph = e4a.apm_graph
-        #inputs = {"question": question, "chat_memory":chat_memory}
-        inputs = {"messages": [{"role": "user", "content": question}]}
-        
-        #add question to memory
-        chat_memory.append(ChatMessage(role="user", content=question))
+    else:
+        response = await apm_graph.ainvoke({"question": question}, config=config)
+        chat_memory.append(ChatMessage(role="assistant", content=response['generation']))
 
-        partial_message = ""
-        async for event in super_graph.astream_events(input=inputs, config=config, version="v2"):
-        #async for event in super_graph.astream(input=inputs, config=config, subgraphs=True):
-            #   chat_memory.append(ChatMessage(role="assistant", content=str(event)))
-            #   yield chat_memory
-
-            kind = event["event"]
-            tags = event.get("tags", [])
-            name =  event['name']
-            
-            #chat_memory.append(ChatMessage(role="assistant", content=f"Running: {name}"))
-            #yield chat_memory
-
-            if name == "safety_check":
-                #if kind == "on_chain_start":
-                #    chat_memory.append(ChatMessage(role="assistant", content=f"- `{name}`"))
-                #    yield chat_memory
-                if kind == "on_chain_stream":
-                    chunk = event['data'].get('chunk')
-                    if chunk and 'safety_status' in chunk and len(chunk['safety_status']) > 0:
-                        chat_memory.append(ChatMessage(role="assistant", content=f"- `{name}`: {chunk['safety_status'][0]}"))
-                        if chunk['safety_status'][0] == 'no' and len(chunk['safety_status']) > 1:
-                            chat_memory.append(ChatMessage(role="assistant", content=f"Safety-status: {chunk['safety_status'][1]}"))
-                    yield chat_memory
-            if kind == "on_chain_end" and name == "route_question":
-                output = event['data'].get('output')
-                if output and 'source' in output:
-                    chat_memory.append(ChatMessage(role="assistant", content=f"- `{name}:` {output['source']}"))
-                else:
-                    chat_memory.append(ChatMessage(role="assistant", content=f"- `{name}:` (no source available)"))
-                yield chat_memory
-            if kind == "on_chain_start" and name == "retrieve":
-                chat_memory.append(ChatMessage(role="assistant", content=f"- `{name}` RAG\n\n"))
-                yield chat_memory
-            if kind == "on_chain_start" and name in ("generate_web_search", "websearch", "stream_generation"):
-                chat_memory.append(ChatMessage(role="assistant", content= f"\n\n- `{name}`\n\n"))
-                yield chat_memory
-            if kind == "on_chain_stream" and name == "stream_generation":
-                data = event["data"]
-                # Accumulate the chunk of data
-                partial_message += data.get('chunk', '')
-                chat_memory[-1].content = partial_message
-                time.sleep(0.05)
-                yield chat_memory
-            if name == "grade_generation_v_documents_and_question":
-                if kind == "on_chain_start": 
-                    chat_memory.append(ChatMessage(role="assistant", content=f"\n\n- `{name}`: "))
-                    yield chat_memory
-                if kind == "on_chain_end":
-                    input_data = event['data'].get('input')
-                    if input_data and hasattr(input_data, 'source'):
-                        output_value = event['data'].get('output', '')
-                        chat_memory.append(ChatMessage(role="assistant", content=f"`{input_data.source}:` {output_value}"))
-                    else:
-                        chat_memory.append(ChatMessage(role="assistant", content=f"`{event['data'].get('output', '')}`"))
-                    yield chat_memory
-            if "stream_hallucination" in tags and kind == "on_chain_start":
-                chat_memory.append(ChatMessage(role="assistant", content=f"- `{tags[-1]}`"))
-                yield chat_memory
-            if "stream_grade_answer" in tags and kind == "on_chain_start":
-                chat_memory.append(ChatMessage(role="assistant", content=f"- `{tags[-1]}`"))
-                yield chat_memory
-            if name == "supervisor":
-                if kind == "on_chain_start":
-                    chat_memory.append(ChatMessage(role="assistant", content=f"- `{name}` "))
-                    yield chat_memory
-                if kind == "on_chain_stream":
-                    chunk = event['data'].get('chunk')
-                    if chunk is not None:
-                        chat_memory.append(ChatMessage(role="assistant", content=f"{chunk}"))
-                        yield chat_memory
-
-            # Set environment variable only when 'event' is defined
-            #os.environ["EA4ALL_" + ea4all_user.replace(".", "_")] = str(event['run_id'])
-        
-        wait_for_all_tracers()
+    yield chat_memory
 
 #Trigger Solution Architecture Diagram QnA
 async def run_vqa_agentic_system(question: str, diagram: str, request: gr.Request) -> AsyncGenerator[list, None]:
@@ -202,39 +116,13 @@ async def run_vqa_agentic_system(question: str, diagram: str, request: gr.Reques
             with Image.open(diagram) as diagram_:
                 if diagram_.format not in allowed_file_types:
                     chat_memory.append(ChatMessage(role="assistant", content="Invalid file type. Allowed file types are JPEG and PNG."))
-                    yield chat_memory
                 else:
                     #'vqa_image = e4u.get_raw_image(diagram) #MOVED into Graph
                     vqa_image = diagram
-
-                    #Setup Quality Assurance Agentic System
-                    #graph = e4v.ea4all_graph(config['configurable']['vqa_model'])
-
-                    #Setup enter graph
-                    diagram_graph = e4v.diagram_graph
-
-                    partial_message = ""
-                    chat_memory.append(ChatMessage(role="assistant", content="Hi, I am working on your question..."))
-                    async for event in diagram_graph.astream_events(
-                        {"question":msg, "image": vqa_image}, config, version="v2"
-                    ):
-                        if (
-                                event["event"] == "on_chat_model_stream"
-                                and "vqa_stream" in event.get('tags', [])
-                                #and event["metadata"].get("langgraph_node") == "tools"
-                            ):
-                                chunk = event["data"].get("chunk")
-                                if chunk is not None and hasattr(chunk, "content"):
-                                    partial_message += chunk.content
-                                    chat_memory[-1].content = partial_message
-                                time.sleep(e4u.CFG.STREAM_SLEEP)
-                                yield chat_memory #, message to update question
-                        elif not partial_message:
-                            yield chat_memory #, message
-                    
-                        #os.environ["EA4ALL_" + ea4all_user.replace(".", "_")] = str(event['run_id'])
-                    
-                    wait_for_all_tracers()
+                    response = await diagram_graph.ainvoke({"question":msg, "image": vqa_image}, config)
+                    chat_memory.append(ChatMessage(role="assistant", content=response['messages'][-1].content))
+                
+                yield chat_memory
 
         except Exception as e:
             yield (e.args[-1])

ea4all/src/ea4all_vqa/graph.py

@@ -20,7 +20,8 @@ from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import (
     AIMessage,
     HumanMessage,
-    ToolMessage
+    ToolMessage,
+    BaseMessage
 )
 
 #pydantic
@@ -73,7 +74,7 @@ class DiagramV2S(BaseModel):
     isSafe: bool = Field(...,description="Should be True if image or question are safe to be processed, False otherwise")
     description: str = Field(description="Should be a string describing the image title.")
 
-@tool("vqa_diagram")
+@tool("vqa_diagram", response_format="content")
 @spaces.GPU
 async def vqa_diagram(next:str, state: Annotated[OverallState, InjectedState], config: RunnableConfig):
     """Diagram Vision Question Answering"""
@@ -100,9 +101,8 @@ async def vqa_diagram(next:str, state: Annotated[OverallState, InjectedState], c
                 },
             ],
         )
-
     prompt = ChatPromptTemplate.from_messages([user_message])
-    values = {"question:":question}
+    values = {"question": question}
 
     llm.max_tokens = set_max_new_tokens(get_predicted_num_tokens_from_prompt(llm, prompt, values))
     chain = prompt | llm
@@ -113,7 +113,8 @@ async def vqa_diagram(next:str, state: Annotated[OverallState, InjectedState], c
 
     response = await chain.ainvoke(input=values, config={"tags": ["vqa_stream"]}, kwargs={"max_tokens": configuration.vqa_max_tokens})
 
-    return response
+    ## When exposed as MCP tool, output schema should as simple as possible as output is serialized to a single string
+    return response.content
 
 ##Supervisor Agent Function custom parse with tool calling response support
 def parse(output: ToolMessage) -> dict | AgentFinish:

ea4all/src/ea4all_vqa/state.py

@@ -15,7 +15,6 @@ from typing import (
 from langchain_core.messages import (
     BaseMessage,
 )
-
 from langgraph.graph import MessagesState
 
 # Optional, the InputState is a restricted version of the State that is used to
@@ -44,21 +43,22 @@ class OutputState:
     """Represents the output schema for the Diagram agent.
     """
 
-    answer: str
-    """Answer to user's question about the Architectural Diagram."""
+    messages: Optional[Annotated[Sequence[MessagesState], operator.add]] = None
+    safety_status: Optional[bool] = None
+
+    """Attributes:
+    safety_status: safety status of the diagram provided by the user
+    Answer to user's question about the Architectural Diagram.
+    """
 
 @dataclass(kw_only=True)
-class OverallState(InputState):
+class OverallState(InputState, OutputState):
     """Represents the overall state of the Diagram graph."""
 
     """Attributes:
-    messages: list of messages
-    safety_status: safety status of the diagram provided by the user
     error: tool error
     next: next tool to be called
     """
     
-    messages: Optional[Annotated[Sequence[BaseMessage], operator.add]] = None
-    safety_status: Optional[bool] = None
     error: Optional[str] = None
     next: Optional[str] = None