Function ai_query

src/codegraphcontext/viz/server.py:560–695

(request: AIQueryRequest)

@app.post("/api/ai_query")
async def ai_query(request: AIQueryRequest):
    if not db_manager:
        raise HTTPException(status_code=500, detail="Database not initialized")
        
    gemini_key = os.getenv("GEMINI_API_KEY")
    openai_key = os.getenv("OPENAI_API_KEY")
    atlascloud_key = os.getenv("ATLASCLOUD_API_KEY") or os.getenv("ATLAS_CLOUD_API_KEY")
    
    if not gemini_key and not openai_key and not atlascloud_key:
        raise HTTPException(
            status_code=400,
            detail=(
                "AI querying requires GEMINI_API_KEY, OPENAI_API_KEY, or ATLASCLOUD_API_KEY. "
                "Please set one in your configuration or .env."
            )
        )
        
    query = request.query
    repo_path = request.repo_path
    
    translation_prompt = """You are an expert graph database administrator. Translate the user's natural language question about a codebase into a valid read-only Cypher query.

The graph database represents the codebase structure with the following schema:
### Nodes & Properties:
* **`Repository`**: `name` (string), `path` (string, absolute path), `is_dependency` (boolean)
* **`File`**: `name` (string), `path` (string, absolute path), `relative_path` (string), `is_dependency` (boolean)
* **`Function`**: `name` (string), `path` (string, absolute path), `line_number` (int), `end_line` (int), `args` (list of strings), `cyclomatic_complexity` (int), `decorators` (list of strings), `lang` (string), `source` (string), `is_dependency` (boolean)
* **`Class`**: `name` (string), `path` (string, absolute path), `line_number` (int), `end_line` (int), `bases` (list of strings), `decorators` (list of strings), `lang` (string), `source` (string), `is_dependency` (boolean)
* **`Interface`**, **`Struct`**, **`Enum`**, **`Trait`**: similar properties to Class/Function
* **`Module`**: `name` (string)

### Relationships:
* **`CONTAINS`**: (Repository)-[:CONTAINS]->(File), (File)-[:CONTAINS]->(Function), (File)-[:CONTAINS]->(Class)
* **`CALLS`**: (Function)-[:CALLS]->(Function)
* **`IMPORTS`**: (File)-[:IMPORTS]->(Module)
* **`INHERITS`**: (Class)-[:INHERITS]->(Class)
* **`IMPLEMENTS`**: (Class)-[:IMPLEMENTS]->(Interface)

### Query Writing Rules:
1. Only write SELECT/MATCH read-only queries. Do NOT use CREATE, MERGE, DELETE, SET, REMOVE, or other modifying keywords.
2. Return nodes using the variable names `n` and `m`, and relationships using `rel` where possible. This is critical for the visualizer.
   Example: MATCH (n:Class)-[rel:INHERITS]->(m:Class) RETURN n, rel, m
3. Keep the query simple and efficient. Limit results using `LIMIT 150` unless more are needed.
4. If the user asks about dependencies or paths, try to return paths using variables, and return nodes/relationships.
5. Make property matches case-insensitive or fuzzy if appropriate (e.g. using `toLower(n.name) CONTAINS 'auth'` or `n.name =~ '(?i).*auth.*'`).
6. If the user is querying a specific repository, you can filter by `Repository.path` or `Repository.name` if path is supplied.

Respond ONLY with a JSON object in this format:
{
  "cypher_query": "MATCH ... RETURN n, rel, m LIMIT 100",
  "explanation": "Brief description of what this Cypher query does"
}
"""
    
    try:
        # Call LLM to get Cypher
        llm_response = call_llm(translation_prompt, query)
        parsed_response = parse_json_from_llm(llm_response)