48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
import dspy
|
|
from typing import List
|
|
from config_loader import load_config
|
|
|
|
CFG = load_config()
|
|
INGESTION_CONFIG = CFG["ingestion_agent"]
|
|
|
|
|
|
class IngestionSignature(dspy.Signature):
|
|
f"{INGESTION_CONFIG['ingestion_signature']}"
|
|
|
|
note: str = dspy.InputField(desc="The DM notes or session recap content.")
|
|
answer: dict[str, str | List] = dspy.OutputField(
|
|
desc="the metadata dictionary with the keys; synopsis, tags, entities, relationships"
|
|
)
|
|
|
|
|
|
class IngestionAgent(dspy.Module):
|
|
def __init__(self):
|
|
self.ingest = dspy.Predict(IngestionSignature)
|
|
|
|
def ingest_with_relationships(self, note: str) -> dict:
|
|
"""Ingest notes and return metadata including extracted relationships."""
|
|
response = self.ingest(note=note)
|
|
result = response.answer
|
|
|
|
if not isinstance(result, dict):
|
|
result = {
|
|
"synopsis": "Failed to parse",
|
|
"tags": [],
|
|
"entities": [],
|
|
"relationships": [],
|
|
}
|
|
|
|
if "relationships" not in result:
|
|
entities = result.get("entities", [])
|
|
relationships = []
|
|
|
|
for i, ent1 in enumerate(entities):
|
|
for ent2 in entities[i + 1 :]:
|
|
relationships.append(
|
|
{"entity1": ent1, "entity2": ent2, "type": "co-occurs_with", "strength": 1}
|
|
)
|
|
|
|
result["relationships"] = relationships
|
|
|
|
return result
|