@@ -153,6 +153,7 @@ async def run_analysis_workflow(
153153 compute_agent_tasks ,
154154 run_domain_agents_parallel ,
155155 )
156+ from app .agents .kg_builder import run_kg_builder
156157 from app .agents .orchestrator import run_orchestrator
157158 from app .agents .strategy import run_strategy
158159 from app .agents .triage import run_triage
@@ -167,7 +168,6 @@ async def run_analysis_workflow(
167168 save_findings_from_output ,
168169 update_finding_entity_ids ,
169170 )
170- from app .services .kg_builder import build_knowledge_graph
171171
172172 settings = get_settings ()
173173 session_factory = _get_sessionmaker ()
@@ -965,7 +965,7 @@ async def run_analysis_workflow(
965965
966966 await db .commit () # Commit findings before KG building
967967
968- # ---- Stage 7: Build Knowledge Graph ----
968+ # ---- Stage 7: Build Knowledge Graph (LLM-based KG Builder) ----
969969 logger .info (
970970 "Pipeline starting stage=kg_builder case=%s workflow=%s" ,
971971 case_id ,
@@ -982,24 +982,60 @@ async def run_analysis_workflow(
982982 )
983983 )
984984
985- (
986- kg_entities_created ,
987- kg_relationships_created ,
988- kg_exact_merges ,
989- ) = await build_knowledge_graph (
985+ kg_builder_task_id = str (uuid4 ())
986+ await emit_agent_started (
990987 case_id = case_id ,
991- workflow_id = workflow_id ,
992- domain_results = domain_results ,
993- db = db ,
988+ agent_type = "kg_builder" ,
989+ task_id = kg_builder_task_id ,
990+ file_id = "" ,
991+ file_name = "knowledge-graph-builder" ,
994992 )
993+
994+ try :
995+ kg_entities_created , kg_relationships_created = await run_kg_builder (
996+ case_id = case_id ,
997+ workflow_id = workflow_id ,
998+ user_id = user_id ,
999+ domain_results = domain_results ,
1000+ db_session = db ,
1001+ publish_event = publish_fn ,
1002+ )
1003+ await emit_agent_complete (
1004+ case_id = case_id ,
1005+ agent_type = "kg_builder" ,
1006+ task_id = kg_builder_task_id ,
1007+ result = {
1008+ "taskId" : kg_builder_task_id ,
1009+ "agentType" : "kg_builder" ,
1010+ "outputs" : [
1011+ {
1012+ "type" : "kg-builder-results" ,
1013+ "data" : {
1014+ "entitiesCreated" : kg_entities_created ,
1015+ "relationshipsCreated" : kg_relationships_created ,
1016+ },
1017+ }
1018+ ],
1019+ },
1020+ )
1021+ except Exception as exc :
1022+ logger .exception ("KG Builder failed for case=%s: %s" , case_id , exc )
1023+ kg_entities_created = 0
1024+ kg_relationships_created = 0
1025+ await emit_agent_error (
1026+ case_id = case_id ,
1027+ agent_type = "kg_builder" ,
1028+ task_id = kg_builder_task_id ,
1029+ error = str (exc )[:500 ],
1030+ )
1031+
9951032 await db .commit () # Commit KG data
9961033
9971034 logger .info (
998- "KG build complete case=%s entities=%d relationships=%d merges=%d " ,
1035+ "KG build complete case=%s entities=%d relationships=%d" ,
9991036 case_id ,
10001037 kg_entities_created ,
10011038 kg_relationships_created ,
1002- kg_exact_merges ,
10031039 )
10041040
10051041 # ---- Stage 7b: Backfill finding-to-entity links ----
0 commit comments