Skip to content

Embed the question once per chat turn, not on every graph lookup - #971

Merged
WaylandYang merged 2 commits into
deeplethe:devfrom
Maya-Kid:fix/embed-once-per-turn
Sep 27, 2026
Merged

WaylandYang merged 2 commits into
deeplethe:devfrom
Maya-Kid:fix/embed-once-per-turn

Conversation

@Maya-Kid

Copy link
Copy Markdown
Contributor

Problem

When a name matches more than one entity, the graph tools rank the candidates by the question. rank_by_question embeds the turn's question and compares it with each candidate's context profile.

The question is the same for the whole turn, but it is embedded again on every lookup by name:

  • find_entities;
  • every entity_facts, neighbors and timeline call that takes a name;
  • paths_between, once for each end.

Each time, ToolCtx::embed reads the workspace's model settings and calls the embedding service.

A predicate that matches no fact is embedded the same way, on every call with that word (aligned_predicates).

Change

  • ToolSink gains embeddings: the texts embedded in this turn, with their vectors.
    • Chat has one sink per turn (Shared.sink) and MCP one per call, so nothing outlives a turn.
    • Only successes are kept: after a failed call, the next lookup tries again, as today.
  • ToolCtx::embed_once(sink, text) returns the kept vector, or calls embed and keeps the result.
  • rank_by_question and aligned_predicates use it. They and filtered now take the sink, which every caller already holds.
  • Rankings don't change. The vectors are identical, so the ranking and the predicate alignment come out the same.

Not in scope, and they could share this later:

  • search_chunks embeds its own query through retrieval;
  • with data sources mounted, mapping_index::relevant also embeds the question.

Verification

Run in containers on a Linux host: Rust 1.98.1 and Postgres 16 with pgvector.

New tests

chat_embed_once_tests.rs uses a fake /embeddings endpoint that records every request. The workspace's embedding model points at it.

The seed is two entities named Acme:

  • the one whose profile is close to the question has fewer facts;
  • ranked by fact count alone, as over MCP, the other one comes first.

Tools, with one sink standing for one turn:

  • With no question, find_entities makes no request and lists the busier Acme first.
  • With the question:
    • find_entities, entity_facts by name and neighbors by name all pick the closer Acme;
    • the endpoint gets one request, for the question;
    • entity_facts twice with a predicate that matches no fact adds one request, for the word.
  • Next turn. A new sink embeds the question again.

Chat: a scripted model calls find_entities, then entity_facts by name. The endpoint gets one request.

On dev (f1a36a5) with only the new tests

Both tests fail on their request counts. The ranking checks before those counts pass on dev too: dev picks the same entities, it only embeds more often.

A copy of the tests that prints the requests instead of asserting them gives the full count:

after requests on dev on this branch
chat turn: find_entities, then entity_facts by name 2 1
tools: find_entities, entity_facts and neighbors by name 3 1
… then entity_facts twice with a predicate that matches no fact 7 2
… then find_entities in the next turn 8 3

On dev, every lookup by name re-embeds the question. The two predicate reads cost two requests each, one for the question and one for the word.

This branch

  • The new tests pass.
  • cargo fmt --all --check is clean.
  • cargo clippy --locked --workspace --all-targets -- -D warnings is clean.
  • cargo test --locked -p utopia-server --bin utopia-server api:: passes 189 tests.
  • UTOPIA_TEST_REQUIRE_PDFTOTEXT=1 cargo test --locked --workspace passes.

The branch sits on f1a36a5 and merges cleanly onto today's dev, which has #967 and #968.

Signed-off-by: dada-yan <BinjunYann@gmail.com>

@WaylandYang WaylandYang left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you. The cache lives on the sink, so it cannot outlive a turn, only successes are kept, and the vectors are the same so nothing ranks differently. The request-count table against dev is the proof. Landing it. search_chunks and mapping_index::relevant sharing the same cache would be a fine follow-up. Thanks again.

@WaylandYang WaylandYang left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maintainer review: sound, landing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants