Configure wrapping

test_langgraph/__main__.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
 from fastapi.encoders import jsonable_encoder
+import uvicorn
 
 from . import tools
 
@@ -26,7 +27,22 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-for tool in tools.get_tools():
-    component, method = tool.__name__.split('.')
-    path = f'/{component}/{method}'
-    app.get(path, response_model=None,tags=[component])(tool)
+def setup_mcp():
+    try:
+        from fastapi_mcp import FastApiMCP
+        mcp = FastApiMCP(app)
+        mcp.mount()
+        print('MCP mounted')
+    except ImportError:
+        print('MCP could not be mounted')
+
+def setup_tools():
+    for tool in tools.get_tools():
+        component, method = tool.__name__.split('.')
+        path = f'/{component}/{method}'
+        app.get(path, response_model=None,tags=[component])(tool)
+
+if __name__ == "__main__":
+    setup_tools()
+    setup_mcp()
+    uvicorn.run(app, host="0.0.0.0", port=8000)

test_langgraph/tools.py

@@ -60,10 +60,11 @@ RETURN_FORMATS = {
 
 RETURN_FORMAT = 'json'
 
-MAX_TOOL_RESULT_LEN = 1000
+
 APPEND_RESULT_TYPE_DOCS = True
 
-def wrap_method(class_, method):
+def wrap_method(class_, method, max_tool_result_len = 1_000_000_000,
+                provide_instructions: bool = False):
     logger.info('Wrapping %s.%s', class_.__name__, method.__name__)
     return_type = method.__annotations__.get('return', '')
     is_iterator = str(return_type).startswith(
@@ -91,10 +92,13 @@ def wrap_method(class_, method):
                 repr(input_value),
             )
             raise
-        if len(result_str) > MAX_TOOL_RESULT_LEN:
-            result_str = result_str[:MAX_TOOL_RESULT_LEN] + ' (remaining tool result elicited...)'
-        if APPEND_RESULT_TYPE_DOCS and (return_docs := getattr(return_type, '__doc__', None)):
-            result_str = result_str+'\n'+return_docs
+        if len(result_str) > max_tool_result_len:
+            result_str = result_str[:max_tool_result_len]
+            if provide_instructions:
+                result_str += '\nremaining tool result elicited...'
+        if provide_instructions:
+            if APPEND_RESULT_TYPE_DOCS and (return_docs := getattr(return_type, '__doc__', None)):
+                result_str += '\n'+return_docs
         return result_str
 
     wrapper.__name__ = f'{class_.__name__}.{method.__name__}'