query_pipeline/thesis_extractor.py

from domain.chunk_d import ChunkD
from domain.entity_d import EntityKnowledgeGraphD
from extraction_pipeline.relationship_extractor.openai_relationship_extractor import (
    OpenAIRelationshipExtractor,)


class ThesisExtractor:
    _relationship_extractor: OpenAIRelationshipExtractor

    def __init__(self):
        self._relationship_extractor = OpenAIRelationshipExtractor()

    def extract_relationships(self, thesis: ChunkD) -> EntityKnowledgeGraphD:
        # Thesis chunk will have DocumentD for its parent_reference field
        # which will contain "user" as the author and more importantly
        # the time of the thesis query for use by the relationship extractor
        entity_relationships = list(self._relationship_extractor.process_element(thesis))
        return EntityKnowledgeGraphD(entity_relationships=entity_relationships)