fix chunks model for knowledge graph (#550)

pingcap · Dec 27, 2024 · ebb31e9 · ebb31e9
1 parent a01d963
commit ebb31e9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/backend/app/rag/knowledge_base/index_store.py b/backend/app/rag/knowledge_base/index_store.py
@@ -29,13 +29,15 @@ def get_kb_tidb_graph_store(session: Session, kb: KnowledgeBase) -> TiDBGraphSto
     entity_model = get_kb_entity_model(kb)
     relationship_model = get_kb_relationship_model(kb)
     inspection.inspect(relationship_model)
+    chunk_model = get_kb_chunk_model(kb)
 
     graph_store = TiDBGraphStore(
         dspy_lm=dspy_lm,
         session=session,
         embed_model=embed_model,
         entity_db_model=entity_model,
         relationship_db_model=relationship_model,
+        chunk_db_model=chunk_model,
     )
     return graph_store