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

ea4all/src/ea4all_apm/graph.py

@@ -516,7 +516,7 @@ def decide_to_generate(state):
         return "generate"
 
 def grade_generation_v_documents_and_question(
-    state:OverallState, config: RunnableConfig):
+    state:OverallState, config: RunnableConfig) -> str:
     """
     Determines whether the generation is grounded in the document and answers question.
 
@@ -560,13 +560,13 @@ def grade_generation_v_documents_and_question(
         grade = output["score"]
         if grade == "yes":
             print("---DECISION: GENERATION ADDRESSES QUESTION---")
-            yield "useful"
+            return "useful"
         else:
             print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
-            yield "not useful"
+            return "not useful"
     else:
         print("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
-        yield "not supported"
+        return "not supported"
 
 async def apm_query_router(
     state: OverallState, config: RunnableConfig

ea4all/src/ea4all_apm/state.py

@@ -33,33 +33,28 @@ class InputState:
 
 @dataclass(kw_only=True)
 class OutputState:
-    """Represents the output schema for the APM agent.
-    """
-
-    answer: str
+    """Represents the output schema for the APM agent."""
+    question: str
+    documents: Optional[List[str]] = None
+    generation: Optional[str] = None
+    web_search: Optional[str] = None
     """Answer to user's Architecture IT Landscape question about ."""
 
 @dataclass(kw_only=True)
-class OverallState(InputState):
+class OverallState(InputState, OutputState):
     """State of the APM graph / agent."""
     
     """
     safety_status: user question's safeguarding status, justification, rephrased question
     router: classification of the user's query
     source: RAG or websearch
-    web_search: whether to add search
     retrieved: list of documents retrieved by the retriever
     rag: last RAG approach used
     chat_memory: user chat memory
-    generation: should the agent generate a response
-    documents: list of documents retrieved by the retriever
     """
     safety_status: Optional[Tuple[str, str, str]] = None
     router: Optional[Router] = None
     source: Optional[str] = None
     rag: Optional[str] = None
-    web_search: Optional[str] = None
     chat_memory: Optional[str] = None
-    generation: Optional[str] = None
     retrieved: Optional[List[str]] = None
-    documents: Optional[List[str]] = None