minor refactoring and cleanup

06_gpu_and_ml/langchains/codelangchain/app.py

@@ -1,6 +1,6 @@
 import agent
 import modal
-from agent import stub
+from agent import nodes, stub
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
@@ -28,20 +28,20 @@ def serve():
     from langchain_core.runnables import RunnableLambda
     from langserve import add_routes
 
-    def inp(question: str):
+    def inp(question: str) -> dict:
         return {"keys": {"question": question, "iterations": 0}}
 
     def out(state: dict) -> str:
-        print(state)
-        if "generate" in state:  # patch issue with langserve playground
-            return "\n".join(
-                str(elem) for elem in state["generate"]["keys"]["generation"]
-            )
-        return state["keys"]["response"]
+        if "keys" in state:
+            return state["keys"]["response"]
+        elif "generate" in state:
+            return nodes.extract_response(state["generate"])
+        else:
+            return str(state)
 
-    runnable = agent.construct_graph(debug=False).compile()
+    graph = agent.construct_graph(debug=False).compile()
 
-    chain = RunnableLambda(inp) | runnable | RunnableLambda(out)
+    chain = RunnableLambda(inp) | graph | RunnableLambda(out)
 
     add_routes(
         web_app,

06_gpu_and_ml/langchains/codelangchain/common.py

@@ -4,7 +4,7 @@
 import modal
 
 image = modal.Image.debian_slim(python_version="3.11").pip_install(
-    "beautifulsoup4~=4.12.3",  # TODO: move to a "retrieval_image"
+    "beautifulsoup4~=4.12.3",
     "langchain==0.1.11",
     "langgraph==0.0.26",
     "langchain_community==0.0.27",

06_gpu_and_ml/langchains/codelangchain/edges.py

@@ -81,7 +81,6 @@ def decide_to_finish(state: GraphState) -> str:
     iter = state_dict["iterations"]
 
     if error == "None" or iter >= 3:
-        # give up :<
         print("---DECISION: FINISH---")
         return "finish"
     else:

06_gpu_and_ml/langchains/codelangchain/nodes.py

@@ -283,12 +283,27 @@ def finish(self, state: GraphState) -> dict:
         """
 
         print("---FINISHING---")
-        state_dict = state["keys"]
-        code_solution = state_dict["generation"]
-        prefix = code_solution[0].prefix
-        imports = code_solution[0].imports
-        code = code_solution[0].code
 
-        response = "\n".join([prefix, imports, code])
+        response = extract_response(state)
 
         return {"keys": {"response": response}}
+
+
+def extract_response(state: GraphState) -> str:
+    """
+    Extract the response from the graph state
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        str: The response
+    """
+
+    state_dict = state["keys"]
+    code_solution = state_dict["generation"][0]
+    prefix = code_solution.prefix
+    imports = code_solution.imports
+    code = code_solution.code
+
+    return "\n".join([prefix, imports, code])