diff --git a/addNum.py b/addNum.py new file mode 100644 index 0000000..5e28be6 --- /dev/null +++ b/addNum.py @@ -0,0 +1,59 @@ +from flask import Flask, request, jsonify +import uuid + +app = Flask(__name__) + +# A registry of tools the server supports +# Each tool is a Python function taking params and returning result +TOOLS = { + "add_numbers": lambda params: params["a"] + params["b"] +} + +def make_jsonrpc_response(result=None, error=None, id=None): + resp = {"jsonrpc": "2.0", "id": id} + if error is not None: + resp["error"] = error + else: + resp["result"] = result + return resp + +@app.route("/mcp", methods=["POST"]) +def mcp_endpoint(): + """ + Expect JSON-RPC 2.0 requests. Handle tool invocation. + """ + body = request.get_json() + # Basic validation + if body is None or "jsonrpc" not in body or body["jsonrpc"] != "2.0": + return jsonify(make_jsonrpc_response( + error={"code": -32600, "message": "Invalid Request"}, id=body.get("id", None) + )) + + method = body.get("method") + params = body.get("params", {}) + req_id = body.get("id") + + # A method for capability discovery, e.g. list tools + if method == "mcp.discover_tools": + # Return the list of tools and their parameter schemas (simplified) + tools_info = { name: {"params": list(func.__code__.co_varnames)} for name, func in TOOLS.items() } + return jsonify(make_jsonrpc_response(result=tools_info, id=req_id)) + + # If method corresponds to a tool + if method in TOOLS: + try: + tool_fn = TOOLS[method] + result = tool_fn(params) + return jsonify(make_jsonrpc_response(result=result, id=req_id)) + except Exception as e: + return jsonify(make_jsonrpc_response( + error={"code": -32000, "message": f"Tool execution error: {e}"}, id=req_id + )) + + # Method not found + return jsonify(make_jsonrpc_response( + error={"code": -32601, "message": "Method not found"}, id=req_id + )) + +if __name__ == "__main__": + app.run(port=5000) diff --git a/src/codegraphcontext/tools/graph_builder.py b/src/codegraphcontext/tools/graph_builder.py index c858fd4..f069fd7 100644 --- a/src/codegraphcontext/tools/graph_builder.py +++ b/src/codegraphcontext/tools/graph_builder.py @@ -354,21 +354,21 @@ def _create_function_calls(self, session, file_data: Dict, imports_map: dict): """Create CALLS relationships with a unified, prioritized logic flow for all call types.""" caller_file_path = str(Path(file_data['file_path']).resolve()) local_function_names = {func['name'] for func in file_data.get('functions', [])} - local_imports = {imp.get('alias') or imp['name'].split('.')[-1]: imp['name'] + local_imports = {imp.get('alias') or imp['name'].split('.')[-1]: imp['name'] for imp in file_data.get('imports', [])} - + for call in file_data.get('function_calls', []): called_name = call['name'] if called_name in __builtins__: continue resolved_path = None - + if call.get('inferred_obj_type'): obj_type = call['inferred_obj_type'] possible_paths = imports_map.get(obj_type, []) if len(possible_paths) > 0: resolved_path = possible_paths[0] - + else: lookup_name = call['full_name'].split('.')[0] if '.' in call['full_name'] else called_name possible_paths = imports_map.get(lookup_name, []) @@ -383,13 +383,28 @@ def _create_function_calls(self, session, file_data: Dict, imports_map: dict): if full_import_name.replace('.', '/') in path: resolved_path = path break - + if not resolved_path: if called_name in imports_map and imports_map[called_name]: resolved_path = imports_map[called_name][0] else: resolved_path = caller_file_path + # Check if the called_name is a class, and if so, redirect to __init__ + is_class_result = session.run(""" + MATCH (c:Class {name: $called_name, file_path: $resolved_path}) + RETURN count(c) > 0 as is_class + """, called_name=called_name, resolved_path=resolved_path).single() + + if is_class_result and is_class_result['is_class']: + # Find the __init__ method of this class + init_result = session.run(""" + MATCH (c:Class {name: $class_name, file_path: $resolved_path})-[:CONTAINS]->(f:Function {name: '__init__'}) + RETURN f.name as init_name + """, class_name=called_name, resolved_path=resolved_path).single() + if init_result: + called_name = '__init__' + caller_context = call.get('context') if caller_context and len(caller_context) == 3 and caller_context[0] is not None: caller_name, _, caller_line_number = caller_context