218 lines
9.6 KiB
Python
218 lines
9.6 KiB
Python
"""Adaptive RAG mixin module."""
|
|
from __future__ import annotations
|
|
|
|
from .shared import *
|
|
|
|
|
|
class AdaptiveRAGRetrievalGraphMixin:
|
|
async def _retrieve_graph_context(
|
|
self,
|
|
question: str,
|
|
*,
|
|
project_id: Optional[int],
|
|
resolved_entities: Optional[dict[str, Any]] = None,
|
|
) -> list[RetrievedItem]:
|
|
try:
|
|
resolution = resolved_entities or {}
|
|
matched_entities = (
|
|
resolution.get("matched_entities") if isinstance(resolution.get("matched_entities"), list) else []
|
|
)
|
|
anchor_ids_raw = (
|
|
resolution.get("anchor_entity_ids") if isinstance(resolution.get("anchor_entity_ids"), list) else []
|
|
)
|
|
anchor_ids: list[int] = []
|
|
for value in anchor_ids_raw:
|
|
try:
|
|
anchor_ids.append(int(value))
|
|
except Exception:
|
|
continue
|
|
|
|
seed_entities: list[KGEntity] = []
|
|
if matched_entities:
|
|
ids: list[int] = []
|
|
for row in matched_entities:
|
|
if not isinstance(row, dict):
|
|
continue
|
|
try:
|
|
ids.append(int(row.get("entity_id")))
|
|
except Exception:
|
|
continue
|
|
if ids:
|
|
seed_entities = (
|
|
await self.db.execute(
|
|
select(KGEntity).where(
|
|
and_(
|
|
KGEntity.tenant_id == self.tenant_id,
|
|
KGEntity.id.in_(ids),
|
|
)
|
|
)
|
|
)
|
|
).scalars().all()
|
|
|
|
if not seed_entities:
|
|
terms = self._extract_query_terms(question, preferred=self._extract_project_lookup_keyword(question))
|
|
if terms:
|
|
conds = []
|
|
for term in terms[:6]:
|
|
like = f"%{term}%"
|
|
conds.append(KGEntity.display_name.ilike(like))
|
|
conds.append(KGEntity.entity_key.ilike(like))
|
|
seed_entities = (
|
|
await self.db.execute(
|
|
select(KGEntity)
|
|
.where(
|
|
and_(
|
|
KGEntity.tenant_id == self.tenant_id,
|
|
or_(*conds),
|
|
)
|
|
)
|
|
.order_by(KGEntity.last_seen_at.desc(), KGEntity.id.desc())
|
|
.limit(15)
|
|
)
|
|
).scalars().all()
|
|
if not seed_entities and anchor_ids:
|
|
seed_entities = (
|
|
await self.db.execute(
|
|
select(KGEntity).where(
|
|
and_(
|
|
KGEntity.tenant_id == self.tenant_id,
|
|
KGEntity.id.in_(anchor_ids[:10]),
|
|
)
|
|
)
|
|
)
|
|
).scalars().all()
|
|
if not seed_entities:
|
|
return []
|
|
|
|
seed_ids = [int(x.id) for x in seed_entities]
|
|
relation_rows = (
|
|
await self.db.execute(
|
|
select(KGRelation)
|
|
.where(
|
|
and_(
|
|
KGRelation.tenant_id == self.tenant_id,
|
|
or_(
|
|
KGRelation.subject_entity_id.in_(seed_ids),
|
|
KGRelation.object_entity_id.in_(seed_ids),
|
|
),
|
|
)
|
|
)
|
|
.order_by(KGRelation.weight.desc(), KGRelation.id.desc())
|
|
.limit(120)
|
|
)
|
|
).scalars().all()
|
|
|
|
entity_ids: set[int] = set(seed_ids)
|
|
for rel in relation_rows:
|
|
entity_ids.add(int(rel.subject_entity_id))
|
|
entity_ids.add(int(rel.object_entity_id))
|
|
entities = (
|
|
await self.db.execute(
|
|
select(KGEntity).where(
|
|
and_(
|
|
KGEntity.tenant_id == self.tenant_id,
|
|
KGEntity.id.in_(list(entity_ids)),
|
|
)
|
|
)
|
|
)
|
|
).scalars().all()
|
|
entity_map = {int(row.id): row for row in entities}
|
|
|
|
relation_ids = [int(row.id) for row in relation_rows]
|
|
evidence_rows: list[KGEvidence] = []
|
|
if relation_ids or entity_ids:
|
|
evidence_rows = (
|
|
await self.db.execute(
|
|
select(KGEvidence)
|
|
.where(
|
|
and_(
|
|
KGEvidence.tenant_id == self.tenant_id,
|
|
or_(
|
|
KGEvidence.relation_id.in_(relation_ids or [-1]),
|
|
KGEvidence.entity_id.in_(list(entity_ids) or [-1]),
|
|
),
|
|
)
|
|
)
|
|
.order_by(KGEvidence.occurred_at.desc(), KGEvidence.id.desc())
|
|
.limit(220)
|
|
)
|
|
).scalars().all()
|
|
evidence_by_relation: dict[int, str] = {}
|
|
evidence_by_entity: dict[int, str] = {}
|
|
for ev in evidence_rows:
|
|
excerpt = str(ev.content_excerpt or "").strip()
|
|
if not excerpt:
|
|
continue
|
|
if ev.relation_id is not None and int(ev.relation_id) not in evidence_by_relation:
|
|
evidence_by_relation[int(ev.relation_id)] = excerpt[:320]
|
|
if ev.entity_id is not None and int(ev.entity_id) not in evidence_by_entity:
|
|
evidence_by_entity[int(ev.entity_id)] = excerpt[:320]
|
|
|
|
out: list[RetrievedItem] = []
|
|
for rel in relation_rows:
|
|
subject = entity_map.get(int(rel.subject_entity_id))
|
|
target = entity_map.get(int(rel.object_entity_id))
|
|
if subject is None or target is None:
|
|
continue
|
|
subject_payload = self._safe_dict(subject.payload)
|
|
target_payload = self._safe_dict(target.payload)
|
|
if project_id is not None:
|
|
subject_pid = subject_payload.get("project_id")
|
|
target_pid = target_payload.get("project_id")
|
|
if subject_pid is not None and int(subject_pid) != int(project_id):
|
|
if target_pid is not None and int(target_pid) != int(project_id):
|
|
continue
|
|
evidence_text = evidence_by_relation.get(int(rel.id))
|
|
if not evidence_text:
|
|
evidence_text = evidence_by_entity.get(int(subject.id)) or evidence_by_entity.get(int(target.id)) or ""
|
|
content = (
|
|
f"图路径: {str(subject.display_name or subject.entity_key)} --[{str(rel.relation_type)}]--> "
|
|
f"{str(target.display_name or target.entity_key)}\n"
|
|
f"证据: {evidence_text or '无显式证据摘录'}"
|
|
)
|
|
weight = self._to_float(rel.weight, 1.0)
|
|
score = min(0.99, max(0.45, 0.62 + min(weight, 1.5) * 0.18))
|
|
out.append(
|
|
RetrievedItem(
|
|
source="graph",
|
|
content=content[:1200],
|
|
score=score,
|
|
source_label="Knowledge Graph",
|
|
metadata={
|
|
"relation_id": int(rel.id),
|
|
"relation_type": str(rel.relation_type or ""),
|
|
"subject_entity_id": int(subject.id),
|
|
"object_entity_id": int(target.id),
|
|
"subject_name": str(subject.display_name or subject.entity_key),
|
|
"object_name": str(target.display_name or target.entity_key),
|
|
"path": f"{subject.id}->{str(rel.relation_type)}->{target.id}",
|
|
},
|
|
)
|
|
)
|
|
|
|
if not out:
|
|
for seed in seed_entities[:DEFAULT_TOP_K]:
|
|
payload = self._safe_dict(seed.payload)
|
|
if project_id is not None and payload.get("project_id") is not None:
|
|
if int(payload.get("project_id")) != int(project_id):
|
|
continue
|
|
excerpt = evidence_by_entity.get(int(seed.id)) or ""
|
|
out.append(
|
|
RetrievedItem(
|
|
source="graph",
|
|
content=f"实体: {str(seed.display_name or seed.entity_key)}\n证据: {excerpt or '无显式证据摘录'}",
|
|
score=0.58,
|
|
source_label="Knowledge Graph",
|
|
metadata={
|
|
"entity_id": int(seed.id),
|
|
"entity_type": str(seed.entity_type or ""),
|
|
"entity_key": str(seed.entity_key or ""),
|
|
},
|
|
)
|
|
)
|
|
out.sort(key=lambda x: x.score, reverse=True)
|
|
return out[: max(DEFAULT_TOP_K, 6)]
|
|
except Exception:
|
|
return []
|
|
|