ea4all-gradio-agents-mcp-hackathon-tools-refactoring-apm-websearch-state

ea4all/ea4all_mcp.py

@@ -1,4 +1,4 @@
-#2025-06-04
+#CHANGELOG: 2025-06-04
 ## Gradio Agents MCP Hackathon: retrofit to expose EA4ALL Agentic System Agents only
 ## Greetings message not working
 ## UI exposing too much tools, need to be refactored

ea4all/src/ea4all_apm/graph.py

@@ -6,6 +6,10 @@ and key functions for processing & routing user queries, generating answer to
 Enterprise Architecture related user questions
 about an IT Landscape or Websearch.
 """
+
+#CHANGELOG: 2025-06-08
+# Refactored to use tools.websearch (changes State, removed web_search)
+
 import os
 
 from langgraph.graph import END, StateGraph
@@ -67,7 +71,7 @@ async def retrieve_documents(
     """
     with vectorstore.make_retriever(config) as retriever:
         response = await retriever.ainvoke(state.question, config)
-        return {"documents": response}
+        return {"messages": response}
 
 async def apm_retriever(config: RunnableConfig):
     with vectorstore.make_retriever(config) as retriever:
@@ -463,12 +467,12 @@ async def grade_documents(state, config: RunnableConfig):
 
     print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
     question = state.question
-    documents = state.documents
+    documents = state.messages
+    source = state.source
     llm = get_llm_client(model=configuration.query_model, api_base_url=configuration.api_base_url)
 
     # Score each doc
     filtered_docs = []
-    web_search = "No"
     for d in documents:
         score = retrieval_grader(llm).ainvoke(
             {"user_question": question, "document": d.page_content}
@@ -483,9 +487,10 @@ async def grade_documents(state, config: RunnableConfig):
             print("---GRADE: DOCUMENT NOT RELEVANT---")
             # We do not include the document in filtered_docs
             # We set a flag to indicate that we want to run web search
-            web_search = "Yes"
+            #web_search = "Yes"
+            source = "websearch"
     
-    return {"documents": filtered_docs, "question": question, "web_search": web_search}
+    return {"documents": filtered_docs, "question": question, "source": source}
 
 def decide_to_generate(state):
     """
@@ -500,10 +505,10 @@ def decide_to_generate(state):
 
     print("---ASSESS GRADED DOCUMENTS---")
     state.question
-    web_search = state.web_search
+    source = state.source
     getattr(state,'documents')
 
-    if web_search == "Yes":
+    if source == "websearch":
         # All documents have been filtered check_relevance
         # We will re-generate a new query
         print(
@@ -530,11 +535,11 @@ def grade_generation_v_documents_and_question(
     configuration = AgentConfiguration.from_runnable_config(config)
 
     question = getattr(state,'question')
-    documents = getattr(state,'documents')
+    documents = getattr(state,'messages')
     generation = getattr(state,'generation')
     llm = get_llm_client(model=configuration.query_model, api_base_url=configuration.api_base_url)
 
-    if getattr(state,'web_search') == "Yes":
+    if getattr(state,'source') == "websearch":
         #print("---CHECK HALLUCINATIONS---")
         hallucination_grader_instance =  hallucination_grader(llm)
         #for output in hallucination_grader_instance.stream(
@@ -654,7 +659,7 @@ async def retrieve(
     
     documents = await final_chain.ainvoke({"user_question": question, "chat_memory":[]})
 
-    return {"documents": format_docs(documents['cdocs']), "question": question, "rag":getattr(state,'rag')}
+    return {"messages": format_docs(documents['cdocs']), "question": question, "rag":getattr(state,'rag')}
 
 ### Edges ###
 def route_to_node(state:OverallState):
@@ -692,17 +697,17 @@ async def stream_generation(
     llm = get_llm_client(model=configuration.query_model, api_base_url=configuration.api_base_url,streaming=configuration.streaming)
 
     documents = None
-    web_search = None
     question = None
+    source = None
     chat_memory = None
     async for s in state:
-        documents = getattr(s,"documents")
-        web_search = getattr(s,"web_search")
+        documents = getattr(s,"messages")
         question = getattr(s,"question")
+        source = getattr(s,"source")
         chat_memory = getattr(s,"chat_memory")
 
     # Prompt Web Search generation
-    if web_search == "Yes":
+    if source == "websearch":
         prompt = PromptTemplate(
             template="""<|begin_of_text|><|start_header_id|>system<|end_header_id|> You are an enterprise architect assistant for question-answering tasks.
             Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. 
@@ -738,9 +743,9 @@ async def generate(
     """
     #print("---GENERATE---")
 
-    documents = getattr(state,'documents')
-    web_search = getattr(state,'web_search') 
-    question = getattr(state,'question')
+    #documents = getattr(state,'messages')[-1].content #documents
+    source = getattr(state,'source') 
+    #question = getattr(state,'question')
 
     ##Triggered by hallucination_grade? 2025-02-21 - NOT USER being edged to END atm
     #2025-02-21: it's being triggered by super_graph supervisor as well - need to review as calling web_search twice
@@ -752,13 +757,14 @@ async def generate(
     #        await retrieve(state, config)
 
     # Generate answer
-    tags = ["websearch_stream"] if web_search == "Yes" else ["apm_stream"]
+    tags = ["websearch_stream"] if source == "websearch" else ["apm_stream"]
     gen = RunnableGenerator(stream_generation).with_config(tags=tags)
     generation=""
     async for message in gen.astream(state):
         generation = ''.join([generation,message])
    
-    return {"documents": documents, "question": question, "generation": generation, "web_search": web_search}
+    #return {"messages": documents.content, "question": question, "generation": generation, "web_search": web_search}
+    return {"generation": generation}
 
 #ea4all-qna-agent-conversational-with-memory
 async def apm_agentic_qna(
@@ -808,7 +814,8 @@ async def apm_agentic_qna(
     
     documents = await final_chain.ainvoke({"user_question": question, "chat_memory":chat_memory})
 
-    return {"documents": format_docs(documents['cdocs']), "question": question, "rag":5, "web_search": "No", "generation": None}
+    #return {"documents": format_docs(documents['cdocs']), "question": question, "rag":5, "generation": None}
+    return {"messages": format_docs(documents['cdocs']), "rag":5}
 
 async def final(state: OverallState):
     return {"safety_status": state}

ea4all/src/ea4all_apm/state.py

@@ -35,9 +35,9 @@ class InputState:
 class OutputState:
     """Represents the output schema for the APM agent."""
     question: str
-    documents: Optional[List[str]] = None
+    messages: Optional[List[str]] = None
     generation: Optional[str] = None
-    web_search: Optional[str] = None
+    source: Optional[str] = None
     """Answer to user's Architecture IT Landscape question about ."""
 
 @dataclass(kw_only=True)
@@ -54,7 +54,6 @@ class OverallState(InputState, OutputState):
     """
     safety_status: Optional[Tuple[str, str, str]] = None
     router: Optional[Router] = None
-    source: Optional[str] = None
     rag: Optional[str] = None
     chat_memory: Optional[str] = None
     retrieved: Optional[List[str]] = None

ea4all/src/ea4all_vqa/graph.py

@@ -21,7 +21,6 @@ from langchain_core.messages import (
     AIMessage,
     HumanMessage,
     ToolMessage,
-    BaseMessage
 )
 
 #pydantic
@@ -44,7 +43,6 @@ from langgraph.graph import (
 )
 from langgraph.prebuilt import ToolNode, tools_condition, InjectedState
 from langgraph.types import Command
-from langgraph.checkpoint.memory import MemorySaver
 
 #import APMGraph packages
 from ea4all.src.ea4all_vqa.configuration import AgentConfiguration
@@ -52,7 +50,6 @@ from ea4all.src.ea4all_vqa.state import InputState, OutputState, OverallState
 
 #import shared packages
 from ea4all.src.shared.configuration import BaseConfiguration
-from ea4all.src.shared.state import State
 from ea4all.src.shared.utils import (
     get_llm_client, 
     _get_formatted_date,
@@ -72,7 +69,7 @@ class DiagramV2S(BaseModel):
     """Check whether the image provided is an architecture diagram or flowchart and safe to be processed."""
     isArchitectureImage: bool = Field(...,description="Should be True if an image is an architecture diagram or flowchart, otherwise False.")
     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.")
+    description: str = Field(description="One sentence describing the reason for being categorised as unsafe or not an architecture image.")
 
 @tool("vqa_diagram", response_format="content")
 @spaces.GPU
@@ -174,9 +171,17 @@ def safeguard_check(state:OverallState, config:RunnableConfig) -> dict:
     question = getattr(state, "question", "Describe the image")
     raw_image = get_raw_image(getattr(state,'image', _join_paths(configuration.ea4all_images,'multi-app-architecture.png')))
 
-    system_message = """You are an expert on identifying images as architecture diagrams, flowchart and whether they are innapropriate content to be processed.
-Given the conversation above, is the image safe to be processed? Does the image represent an architecture diagram or flowchart?"
-"""
+    system_message = (
+    "Act as a safeguarding agent to check whether the image provided is an architecture diagram or flowchart and safe to be processed. "
+    "You will be provided with a question and an image. "   
+    "You should return a JSON object with the following fields: "
+    "'isArchitectureImage':bool, 'isSafe': bool, 'description': str. "
+    "The 'isArchitectureImage' field should be True if the image is an architecture diagram or flowchart, otherwise False. "
+    "The 'isSafe' field should be True if the image or question are safe to be processed, False otherwise. "
+    "The 'description' field should contain a one sentence description of the reason for being categorised as unsafe or not an architecture image. "
+    "If the image is not an architecture diagram or flowchart, you should say it is not an architecture image as 'description' field. "
+    "If the image is not safe to be processed, you should say it is unsafe as 'description' field. "
+    )
 
     safeguard_checker = create_safeguarding_agent(
         llm,
@@ -188,11 +193,6 @@ Given the conversation above, is the image safe to be processed? Does the image
     input = {"question": question, "raw_image": raw_image}
     result = safeguard_checker.invoke(input=input, config=config)
     
-    # Parse out the function call
-    architecture_image = result['isArchitectureImage']
-    safe_request = result['isSafe']
-    description = result['description']
-
     return {"safety_status": result}
 
 def call_finish(state:OverallState, config:RunnableConfig) -> dict:

ea4all/src/tools/tools.py

@@ -64,7 +64,7 @@ def make_supervisor_node(config: RunnableConfig, members: list[str]) -> Runnable
 
     return RunnableLambda(supervisor_node)
 
-async def websearch(state: State) -> dict[str,dict[str,str]]:
+async def websearch(state: dict[str, dict | str]) -> dict[str,dict[str,str]]:
     """
     Web search based on the re-phrased question.
 
@@ -84,7 +84,7 @@ async def websearch(state: State) -> dict[str,dict[str,str]]:
         bing_search_url=bing_search_url
     )
 
-    question = state.get('messages')[-1].content
+    question = getattr(state,'messages')[-1].content if getattr(state,'messages', False) else getattr(state,'question')
 
     ##Bing Search Results
     web_results = BingSearchResults(

ea4all/utils/utils.py

@@ -115,11 +115,11 @@ async def ea4all_agent_init(request:gr.Request):
 
     agentic_qna_desc="""Hi,
         improve effieciency, knowledge sharing, and get valuable insights from your IT landscape using natural language.
-        As an Enterprise Architect Agentic System I can answer questions related to Enterprise Architecture, Technology, plus the following IT Landscape sample dataset: """
+        As an Enterprise Architect Agentic Companion I can answer questions related to Enterprise Architecture, Technology, plus the following IT Landscape sample dataset: """
 
     #capture user IP address
     #ea4all_user = e4u.get_user_identification(request)
-    gr.Info("Thank you for connecting! I'd love to hear your feedback! Thumbs up or Thumbs down. LinkedIn comment.")
+    gr.Info("Thank you for using the EA4ALL Agentic MCP Server!")
 
     # Set initial landscape vectorstore