Skip to content

README.md

Latest commit

 

History

History
50 lines (37 loc) · 3.39 KB

README.md

File metadata and controls

50 lines (37 loc) · 3.39 KB
Raw

Retrieval Pipeline Function Directory

Structured retrieval (Analytics / “Text-to-Pandas”)

For numerical / aggregation questions (e.g., “average tax rate for corporations in 2022”), the system uses a structured query layer instead of free-form Pandas code. This keeps analytics execution safe, predictable, and easy to validate.

  • Natural language → structured query: the analytics agent produces a DFQuery object (see structured.py) rather than emitting executable Python.
  • Strict schema + validation: DFQuery validates requested fields against ALLOWED_COLUMNS and restricts operators/aggregations to a small allow-list.
  • Filter sanitization for categorical fields: categorical filters are normalized through CATEGORICAL_MAPPING so common variants (case/spacing/hyphens/underscores) resolve to canonical values (e.g., non-profit → Non-Profit).
  • Runtime safety + determinism: safe_df_query() clamps limits to MAX_QUERY_LIMIT, applies filters via a controlled operator map, and returns JSON-serializable outputs.

Query support

  • Operations:
    • select: choose one or more columns (drops duplicates and applies deterministic limiting)
    • filter: return filtered rows (limited)
    • aggregate / filtered_aggregate: compute scalar or grouped aggregations
  • Aggregations: count, sum, mean, min, max
  • Grouping: group_by can be a string or list of strings; results are flattened to unique column names like mean_income, max_tax_owed.

Graph Expansion

Graph expansion is the second stage of retrieval after semantic search: you start from one KG node id (often a chunk id) and walk the Neo4j property graph to pull nearby facts as compact, LLM-friendly context. In this repo, the preferred interface is get_kg_node_relations_nlp(), which returns a flattened string of edges like <SUBJECT> --RELATION--> <OBJECT> so you can keep token usage low while still capturing structure.

One wrinkle is that chunk nodes are frequently connected via high-volume MENTIONED_IN edges; these edges are useful for traversal but usually low-signal for output. The implementation therefore allows traversing through MENTIONED_IN as a “bridge” when needed, but does not print MENTIONED_IN lines in the final NLP output, and it tries to fill up to max_edges with non-MENTIONED_IN relationships.

Use a shallow depth for quick, local context and increase depth when you need broader recall, while controlling output size with max_edges and max_edges_per_subject. If a chunk appears to have “only mentioned_in” direct edges, you can still get meaningful relations because the expansion will bridge through those neighbors and then surface their non-MENTIONED_IN edges.

from retrieval.graph import get_kg_node_relations_nlp

kg_id = "5422c5f2-c927-4d29-8f11-2aba6e775a6d"
nlp = get_kg_node_relations_nlp(ids=[kg_id], depth=3, max_edges=25)
print(nlp)

Node ID flow:

Vector Search Result (NodeWithScore)
    ↓
node.node_id = "chroma-doc-id-123"  ← Chroma document ID
node.relationships[NodeRelationship.SOURCE].node_id = "neo4j-uuid-456"  ← Neo4j graph ID
    ↓
scored_candidates["node_id"] = "neo4j-uuid-456"  ← Correct ID for expansion
    ↓
selection_map["N1"] = "neo4j-uuid-456"  ← Used in workflow
    ↓
traversal_expand_node resolves selected_choice_id → chunk_id for graph_expand_triplets_tool