Fix failing chain generation

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)

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