Accept traced entity capture citations and improve correction feedback

tests/test_entity_resolution.py

@@ -55,6 +55,7 @@ class Model:
         self.invalid_assessment = 0
         self.unavailable_stage = None
         self.omit_resolved = 0
+        self.linked_capture = False
 
     async def generate(self, config, request, *, max_response_bytes):
         self.requests.append(copy.deepcopy(request))
@@ -83,7 +84,12 @@ class Model:
             data = json.loads(request["messages"][1]["content"])["input"]
             selected = [(k,v) for k,v in data["sources"].items() if v["kind"] == "entity" and v["case_pointer"].startswith("/entity_resolution/")]
             for k,v in selected:
-                document["relevant_entities"].append({"entity_id": k, "reason": "Contextually matched affected switch", "source_ids": [k,"O0001"]})
+                evidence = k
+                if self.linked_capture:
+                    evidence = next(f for f,s in data['sources'].items() if s['kind']=='fact'
+                        and s['data'].get('name')=='entity_name_lookup'
+                        and s['data']['provenance']['source_ref']==v['data']['provenance']['source_ref'])
+                document["relevant_entities"].append({"entity_id": k, "reason": "Contextually matched affected switch", "source_ids": [evidence,"O0001"]})
             if self.omit_resolved:
                 self.omit_resolved -= 1
                 document["relevant_entities"] = payload()["relevant_entities"]
@@ -118,7 +124,7 @@ def test_resolves_unrecorded_subject_preserves_case_and_portable_replay(built_ca
     model = Model(); service = service_for(store)
     result = run(store, model, service)
     assert result.status == "COMPLETED" and PackageValidator().validate(result.package).valid
-    assert [r["prompt_version"] for r in model.requests] == ["entity_resolution_plan_v2","entity_resolution_assessment_v2","case_interpretation_resolved_v1"]
+    assert [r["prompt_version"] for r in model.requests] == ["entity_resolution_plan_v2","entity_resolution_assessment_v2","case_interpretation_resolved_v2"]
     assert len(service.database.calls) == 1
     assert service.database.calls[0]["filters"] == [{"column":"name","operator":"contains","value":"switch"}]
     with PackageReader(result.package) as reader:
@@ -254,10 +260,55 @@ def test_final_interpretation_cannot_discard_resolved_identity(built_case):
     assert result.status=='REJECTED' and result.code=='INTERPRETATION_RESOLVED_ENTITY_OMITTED'
     result=run(store,model,service,correction_config=CorrectionConfig())
     assert result.status=='COMPLETED' and len(model.requests)==4
+    feedback=json.loads(model.requests[-1]['messages'][-1]['content'])['validation_feedback']
+    assert any(i.get('constraint')=='retain_resolved_entity' and i['entity_id']=='N0002' for i in feedback['issues'])
     assert len(service.database.calls)==1
     assert run(store,model).reused and len(model.requests)==4
 
 
+def test_matching_database_capture_is_valid_without_redundant_entity_citation(built_case):
+    store=built_case[0];model=Model();model.linked_capture=True;service=service_for(store)
+    result=run(store,model,service)
+    assert result.status=='COMPLETED' and len(model.requests)==3
+    assert run(store,model).reused
+    with PackageReader(result.package) as r:
+        doc=json.loads(r.read_bytes(result.interpretation_ref))
+        assert 'N0002' not in doc['interpretation']['relevant_entities'][-1]['source_ids']
+
+
+def test_provenance_feedback_identifies_missing_mention_and_wrong_capture(built_case):
+    from viano_bug_analysis.analysis.contracts import entity_citation_issues,validate_interpretation
+    store=built_case[0];model=Model();service=service_for(store);result=run(store,model,service)
+    with PackageReader(result.package) as r:
+        doc=json.loads(r.read_bytes(result.interpretation_ref));context=json.loads(r.read_bytes(doc['context_ref']))
+    item=doc['interpretation']['relevant_entities'][-1]
+    item['source_ids']=['F0001']
+    issues=entity_citation_issues(doc['interpretation'],context)
+    assert {i['constraint'] for i in issues}=={'cite_original_mention','cite_entity_or_matching_capture'}
+    assert next(i for i in issues if i['constraint']=='cite_original_mention')['allowed_source_ids']==['O0001']
+    assert 'N0002' in next(i for i in issues if i['constraint']=='cite_entity_or_matching_capture')['allowed_source_ids']
+    with pytest.raises(InterpretationError,match='INTERPRETATION_ENTITY_PROVENANCE'):
+        validate_interpretation(doc['interpretation'],context)
+    # Even a lookup-shaped fact cannot substitute an unrelated capture.
+    context['sources']['F0001']['data'].update(name='entity_name_lookup',provenance={'source_ref':'different/capture.json'})
+    item['source_ids']=['O0001','F0001']
+    assert [i['constraint'] for i in entity_citation_issues(doc['interpretation'],context)]==['cite_entity_or_matching_capture']
+
+
+def test_real_gemini_reply_is_accepted_with_matching_capture_and_legacy_rule_is_preserved():
+    from pathlib import Path
+    from viano_bug_analysis.analysis.contracts import validate_interpretation
+    fixture=json.loads((Path(__file__).parent/'fixtures/interpretation/real-capture-citations.json').read_text(encoding='utf-8'))
+    validate_interpretation(fixture['response'],fixture['context'])
+    entity=next(e for e in fixture['response']['relevant_entities'] if e['entity_id']=='N0005')
+    assert 'N0005' not in entity['source_ids'] and 'F0027' in entity['source_ids']
+    assert fixture['context']['sources']['N0005']['data']['id']=='23487'
+    old=copy.deepcopy(fixture['context'])
+    del old['sources']['N0005']['data']['citation_policy']
+    with pytest.raises(InterpretationError,match='INTERPRETATION_RESOLVED_ENTITY_OMITTED'):
+        validate_interpretation(fixture['response'],old)
+
+
 @pytest.mark.parametrize('settings', [{'enabled':True,'catalogs':[]},{**SETTINGS,'max_searches':True},{**SETTINGS,'unexpected':5}])
 def test_invalid_operator_configuration(settings):
     with pytest.raises(InterpretationError,match='ENTITY_RESOLUTION_CONFIG'):policy(settings)