first real commit

app.py

@@ -0,0 +1,61 @@
+import gradio as gr
+from typing import Dict, Union, List
+import json
+import uuid
+from storage.neo4j_dao import Neo4jDomainDAO
+from domain.chunk_d import DocumentD, ChunkD
+from query_pipeline.thesis_extractor import ThesisExtractor
+from query_pipeline.evaluation_engine import EvaluationEngine
+from proto.chunk_pb2 import Chunk, ChunkType
+from datetime import datetime
+import os
+
+os.environ['OPENAI_API_KEY'] = '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'
+
+os.environ['NEO4J_USER'] = "neo4j"
+os.environ['NEO4J_PASSWORD'] = "dOIZwzF_GImgwjF-smChe60QGQgicq8ER8RUlZvornU"
+os.environ['NEO4J_URI'] = "neo4j+s://2317ae21.databases.neo4j.io"
+
+def thesis_evaluation(thesis: str) -> str:
+    thesis_document_d = DocumentD(file_path="", 
+                              authors="user",
+                              publish_date="2024-06-18")
+
+    thesis_chunk_d = ChunkD(chunk_text=thesis, 
+                        chunk_type=ChunkType.CHUNK_TYPE_PAGE,
+                        chunk_index=0,
+                        parent_reference=thesis_document_d)
+    
+    thesis_entity_kg = ThesisExtractor().extract_relationships(thesis_chunk_d)
+
+    with Neo4jDomainDAO() as neo4j_dao:
+        full_db_kg = neo4j_dao.get_knowledge_graph()
+
+    return EvaluationEngine(full_db_kg).evaluate_and_display_thesis(thesis_entity_kg)
+
+
+theme = gr.themes.Soft(
+    primary_hue="cyan",
+    secondary_hue="green",
+    neutral_hue="slate",
+)
+
+with gr.Blocks(theme=theme) as demo:
+    gr.HTML(""" 
+            <div style="background-color: rgb(171, 188, 251);">
+                <div style="background-color: rgb(151, 168, 231); padding-top: 2px; padding-bottom: 7px; text-align: center;">
+                    <h1>AthenaAIC MetisLLM Thesis Evaluation Tool</h1>
+                </div>
+            </div>""")
+    with gr.Row():
+        inp = gr.Textbox(placeholder="What is your thesis?", scale=3, lines=3, show_label=False)
+        submit = gr.Button(scale=1, value="Evaluate")
+    with gr.Row():
+        out = gr.HTML()
+
+    submit.click(thesis_evaluation, inputs=inp, outputs=out)
+
+demo.launch(
+    auth=('athena-admin', 'athena'),
+    share=True
+)

domain/chunk_d.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+import dataclasses
+import uuid
+from typing import Union
+import hashlib
+
+import proto.chunk_pb2 as chunk_pb2
+from domain.domain_protocol import DomainProtocol
+
+
+@dataclasses.dataclass(frozen=True)
+class DocumentD(DomainProtocol[chunk_pb2.Document]):
+    file_path: str
+    authors: str
+    publish_date: str
+
+    @property
+    def id(self) -> str:
+        return hashlib.sha256(self.to_proto().SerializeToString()).hexdigest()
+
+    @classmethod
+    def _from_proto(cls, proto: chunk_pb2.Document) -> DocumentD:
+        return cls(file_path=proto.file_path,
+                   authors=proto.authors,
+                   publish_date=proto.publish_date)
+
+    def to_proto(self) -> chunk_pb2.Document:
+        return chunk_pb2.Document(file_path=self.file_path,
+                                  authors=self.authors,
+                                  publish_date=self.publish_date)
+
+
+@dataclasses.dataclass(frozen=True)
+class ChunkD(DomainProtocol[chunk_pb2.Chunk]):
+
+    @property
+    def id(self) -> str:
+        return str(self.chunk_id)
+
+    chunk_text: str
+    chunk_type: chunk_pb2.ChunkType
+    chunk_index: int
+    parent_reference: Union[uuid.UUID, DocumentD]
+    chunk_id: uuid.UUID = dataclasses.field(default_factory=uuid.uuid4)
+
+    def __post_init__(self):
+        if self.chunk_type == chunk_pb2.ChunkType.CHUNK_TYPE_PAGE:
+            if not isinstance(self.parent_reference, DocumentD):
+                raise ValueError(
+                    f"Chunk (id: {self.chunk_id}) with type {self.chunk_type} must have a DocumentD parent_reference."
+                )
+        elif not isinstance(self.parent_reference, uuid.UUID):
+            raise ValueError(
+                f"Chunk (id: {self.chunk_id}) with type {self.chunk_type} must have a uuid.UUID parent_reference."
+            )
+
+    @classmethod
+    def _from_proto(cls, proto: chunk_pb2.Chunk) -> ChunkD:
+        if proto.HasField('parent_chunk_id'):
+            return cls(chunk_id=uuid.UUID(proto.chunk_id),
+                       parent_reference=uuid.UUID(proto.parent_chunk_id),
+                       chunk_text=proto.chunk_text,
+                       chunk_type=proto.chunk_type,
+                       chunk_index=proto.chunk_index)
+        elif proto.HasField('document'):
+            return cls(chunk_id=uuid.UUID(proto.chunk_id),
+                       parent_reference=DocumentD._from_proto(proto.document),
+                       chunk_text=proto.chunk_text,
+                       chunk_type=proto.chunk_type,
+                       chunk_index=proto.chunk_index)
+        else:
+            raise ValueError(
+                f"Chunk proto (id: {proto.chunk_id}) has no 'parent' or 'document' field.")
+
+    def to_proto(self) -> chunk_pb2.Chunk:
+        chunk_proto = chunk_pb2.Chunk()
+        chunk_proto.chunk_id = str(self.chunk_id)
+        chunk_proto.chunk_text = self.chunk_text
+        chunk_proto.chunk_type = self.chunk_type
+        chunk_proto.chunk_index = self.chunk_index
+        if isinstance(self.parent_reference, uuid.UUID):
+            chunk_proto.parent_chunk_id = str(self.parent_reference)
+        elif isinstance(self.parent_reference, DocumentD):
+            chunk_proto.document.CopyFrom(self.parent_reference.to_proto())
+        else:
+            raise ValueError(
+                f"Chunk (id: {self.chunk_id}) parent_reference is of unknown type: {type(self.parent_reference)}"
+            )
+        return chunk_proto

domain/chunk_d_test.py

@@ -0,0 +1,71 @@
+import unittest
+import logging
+import uuid
+
+from domain.chunk_d import DocumentD, ChunkD
+import proto.chunk_pb2 as chunk_pb2
+
+
+class DocumentDTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.document_d = DocumentD(file_path='test_file',
+                                   authors='BofA John Doe Herb Johnson Taylor Mason',
+                                   publish_date='2021-01-01')
+
+    def test_proto_roundtrip(self):
+        document_proto = self.document_d.to_proto()
+        roundtrip_document_d = DocumentD.from_proto(document_proto)
+        self.assertEqual(self.document_d, roundtrip_document_d)
+
+
+class ChunkDTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.parent_reference_document_d = DocumentD(
+            file_path='test_file',
+            authors='BofA John Doe Herb Johnson Taylor Mason',
+            publish_date='2021-01-01')
+
+    def test_validate(self):
+        with self.assertRaises(ValueError,
+                               msg="CHUNK_TYPE_PAGE must have a DocumentD parent_reference"):
+            ChunkD(chunk_text='test chunk text',
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PAGE,
+                   chunk_index=9,
+                   parent_reference=uuid.uuid4(),
+                   chunk_id=uuid.uuid4())
+        with self.assertRaises(ValueError,
+                               msg="CHUNK_TYPE_SENTENCE must have a uuid.UUID parent_reference"):
+            ChunkD(chunk_text='test chunk text',
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_SENTENCE,
+                   chunk_index=0,
+                   parent_reference=self.parent_reference_document_d,
+                   chunk_id=uuid.uuid4())
+        ChunkD(chunk_text='test chunk text',
+               chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PAGE,
+               chunk_index=9,
+               parent_reference=self.parent_reference_document_d,
+               chunk_id=uuid.uuid4())
+        ChunkD(chunk_text='test chunk text',
+               chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_SENTENCE,
+               chunk_index=0,
+               parent_reference=uuid.uuid4(),
+               chunk_id=uuid.uuid4())
+
+    def test_proto_roundtrip(self):
+        test_chunk_d = ChunkD(chunk_id=uuid.uuid4(),
+                              parent_reference=self.parent_reference_document_d,
+                              chunk_text='test chunk text',
+                              chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PAGE,
+                              chunk_index=9)
+        chunk_proto = test_chunk_d.to_proto()
+        roundtrip_chunk_d = ChunkD.from_proto(chunk_proto)
+        self.assertEqual(test_chunk_d, roundtrip_chunk_d)
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    unittest.main()

domain/domain_protocol.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from typing import Optional, Protocol, Tuple, Type, TypeVar, get_args
+
+from google.protobuf import json_format, message
+
+MessageType = TypeVar("MessageType", bound=message.Message)
+DomainProtocolType = TypeVar("DomainProtocolType", bound='DomainProtocol')
+
+
+class ProtoDeserializationError(Exception):
+    ...
+
+
+class DomainProtocol(Protocol[MessageType]):
+
+    @property
+    def id(self) -> str:
+        ...
+
+    @classmethod
+    def _from_proto(cls: Type[DomainProtocolType], proto: MessageType) -> DomainProtocolType:
+        ...
+
+    def to_proto(self) -> MessageType:
+        ...
+
+    @classmethod
+    def message_cls(cls: Type[DomainProtocolType]) -> Type[MessageType]:
+        orig_bases: Optional[Tuple[Type[MessageType], ...]] = getattr(cls, "__orig_bases__", None)
+        if not orig_bases:
+            raise ValueError(f"Class {cls} does not have __orig_bases__")
+        if len(orig_bases) != 1:
+            raise ValueError(f"Class {cls} has unexpected number of bases: {orig_bases}")
+        return get_args(orig_bases[0])[0]
+
+    @classmethod
+    def from_proto(cls: Type[DomainProtocolType],
+                   proto: MessageType,
+                   allow_empty: bool = False) -> DomainProtocolType:
+        try:
+            if not allow_empty:
+                cls.validate_proto_not_empty(proto)
+            return cls._from_proto(proto)
+        except Exception as e:
+            error_str = f"Failed to convert {cls} - {e}"
+            raise ProtoDeserializationError(error_str) from e
+
+    @classmethod
+    def from_json(cls: Type[DomainProtocolType], json_str: str) -> DomainProtocolType:
+        try:
+            proto_cls = cls.message_cls()
+            proto = proto_cls()
+            json_format.Parse(json_str, proto)
+            return cls.from_proto(proto)
+        except json_format.ParseError as e:
+            error_str = f"{cls} failed to parse json string: {json_str} - {e}"
+            raise ProtoDeserializationError(error_str) from e
+
+    def to_json(self) -> str:
+        return json_format.MessageToJson(self.to_proto()).replace("\n", " ")
+
+    @classmethod
+    def validate_proto_not_empty(cls, proto: message.Message):
+        if cls.is_empty(proto):
+            raise ValueError("Proto is empty")
+
+    @classmethod
+    def is_empty(cls, proto: message.Message) -> bool:
+        descriptor = getattr(proto, 'DESCRIPTOR', None)
+        fields = list(descriptor.fields) if descriptor else []
+        while fields:
+            field = fields.pop()
+
+            if field.label == field.LABEL_REPEATED:
+                eval_func = lambda x: x == field.default_value
+                if field.type == field.TYPE_MESSAGE:
+                    eval_func = cls.is_empty
+                if not all([eval_func(item) for item in getattr(proto, field.name)]):
+                    return False
+
+            elif field.type == field.TYPE_MESSAGE:
+                if not cls.is_empty(getattr(proto, field.name)):
+                    return False
+            else:
+                field_value = getattr(proto, field.name)
+                if field_value != field.default_value:
+                    return False
+        return True

domain/domain_protocol_test.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+import logging
+import unittest
+from google.protobuf import timestamp_pb2
+import dataclasses
+
+from domain.domain_protocol import DomainProtocol, ProtoDeserializationError
+
+
+@dataclasses.dataclass(frozen=True)
+class TimestampTestD(DomainProtocol[timestamp_pb2.Timestamp]):
+
+    nanos: int
+
+    @property
+    def id(self) -> str:
+        return str(self.nanos)
+
+    @classmethod
+    def _from_proto(cls, proto: timestamp_pb2.Timestamp) -> TimestampTestD:
+        return cls(nanos=proto.nanos)
+
+    def to_proto(self) -> timestamp_pb2.Timestamp:
+        return timestamp_pb2.Timestamp(nanos=self.nanos)
+
+
+class DomainProtocolTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.timestamp_d = TimestampTestD(nanos=1)
+        cls.timestamp_proto = timestamp_pb2.Timestamp(nanos=1)
+
+    def test_proto_roundtrip(self):
+        proto = self.timestamp_d.to_proto()
+        domain_from_proto = TimestampTestD.from_proto(proto)
+        self.assertEqual(self.timestamp_d, domain_from_proto)
+
+    def test_json_roundtrip(self):
+        json_str = self.timestamp_d.to_json()
+        domain_from_json = TimestampTestD.from_json(json_str)
+        self.assertEqual(self.timestamp_d, domain_from_json)
+
+    def test_from_proto_empty_fail(self):
+        empty_proto = timestamp_pb2.Timestamp()
+        with self.assertRaises(ProtoDeserializationError):
+            TimestampTestD.from_proto(empty_proto)
+
+    def test_from_proto_empty_allowed_flag(self):
+        empty_proto = timestamp_pb2.Timestamp()
+        domain_from_proto = TimestampTestD.from_proto(empty_proto, allow_empty=True)
+        self.assertEqual(TimestampTestD(nanos=0), domain_from_proto)
+
+    def test_validate_proto_not_empty(self):
+        empty_proto = timestamp_pb2.Timestamp()
+        with self.assertRaises(ValueError):
+            TimestampTestD.validate_proto_not_empty(empty_proto)
+
+    def test_is_empty(self):
+        empty_proto = timestamp_pb2.Timestamp()
+        self.assertTrue(TimestampTestD.is_empty(empty_proto))
+
+    def test_message_cls(self):
+        self.assertEqual(timestamp_pb2.Timestamp, TimestampTestD.message_cls())
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    unittest.main()

domain/entity_d.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import dataclasses
+import hashlib
+import logging
+from typing import TypeAlias, Union
+
+import proto.entity_pb2 as entity_pb2
+
+from domain.domain_protocol import DomainProtocol
+from utils.dates import parse_date
+
+Neo4jDict: TypeAlias = dict[str, Union[str, int, bool, list[str], list[int], list[bool]]]
+
+
+@dataclasses.dataclass(frozen=True)
+class EntityD(DomainProtocol[entity_pb2.Entity]):
+    entity_id: str
+    entity_name: str
+
+    @property
+    def id(self) -> str:
+        return self.entity_id
+
+    @classmethod
+    def _from_proto(cls, proto: entity_pb2.Entity) -> EntityD:
+        return EntityD(entity_id=proto.entity_id, entity_name=proto.entity_name)
+
+    def to_proto(self) -> entity_pb2.Entity:
+        return entity_pb2.Entity(entity_id=self.entity_id, entity_name=self.entity_name)
+
+    @property
+    def neo4j_create_cmd(self):
+        # TODO store entity_id?
+        return "MERGE (e:Entity {name: $name}) ON CREATE SET e.pdf_file = $pdf_file"
+
+    @property
+    def neo4j_create_args(self) -> Neo4jDict:
+        return {
+            "name": self.entity_name,
+        }
+
+
+@dataclasses.dataclass(frozen=True)
+class RelationshipD(DomainProtocol[entity_pb2.Relationship]):
+    relationship_id: str
+    start_date: str
+    end_date: str
+    source_text: str
+    predicted_movement: entity_pb2.PredictedMovement
+
+    @property
+    def id(self) -> str:
+        return self.relationship_id
+
+    def __post_init__(self):
+        if self.start_date and self.end_date:
+            start = parse_date(self.start_date)
+            end = parse_date(self.end_date)
+
+            if end < start:
+                logging.warning("end_date %s is before start_date %s",
+                                self.end_date,
+                                self.start_date)
+                # raise ValueError(f"end_date {self.end_date} is before start_date {self.start_date}")
+
+    @classmethod
+    def _from_proto(cls, proto: entity_pb2.Relationship) -> RelationshipD:
+        return RelationshipD(relationship_id=proto.relationship_id,
+                             start_date=proto.start_date,
+                             end_date=proto.end_date,
+                             source_text=proto.source_text,
+                             predicted_movement=proto.predicted_movement)
+
+    def to_proto(self) -> entity_pb2.Relationship:
+        return entity_pb2.Relationship(relationship_id=self.relationship_id,
+                                       start_date=self.start_date,
+                                       end_date=self.end_date,
+                                       source_text=self.source_text,
+                                       predicted_movement=self.predicted_movement)
+
+    @classmethod
+    def from_string(cls, relationship: str) -> entity_pb2.PredictedMovement:
+        if relationship == "PREDICTED_MOVEMENT_NEUTRAL":
+            return entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL
+        elif relationship == "PREDICTED_MOVEMENT_INCREASE":
+            return entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_INCREASE
+        elif relationship == "PREDICTED_MOVEMENT_DECREASE":
+            return entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_DECREASE
+        else:
+            return entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_UNSPECIFIED
+
+    @property
+    def neo4j_create_cmd(self):
+        return """MATCH (from:Entity {name: $from_name})
+        MATCH (to:Entity {name: $to_name})
+        MERGE (from) -[r:Relationship {start_date: $start_date, end_date: $end_date, predicted_movement: $predicted_movement}]-> (to) ON CREATE SET r.source_text = $source_text, r.pdf_file = $pdf_file"""
+
+    @property
+    def neo4j_create_args(self) -> Neo4jDict:
+        return {
+            "start_date": self.start_date,
+            "end_date": self.end_date,
+            "predicted_movement": entity_pb2.PredictedMovement.Name(self.predicted_movement),
+            "source_text": self.source_text,
+        }
+
+
+@dataclasses.dataclass(frozen=True)
+class EntityRelationshipD(DomainProtocol[entity_pb2.EntityRelationship]):
+    from_entity: EntityD
+    relationship: RelationshipD
+    to_entity: EntityD
+
+    @property
+    def id(self) -> str:
+        return hashlib.sha256(self.to_proto().SerializeToString()).hexdigest()
+
+    @classmethod
+    def _from_proto(cls, proto: entity_pb2.EntityRelationship) -> EntityRelationshipD:
+        return EntityRelationshipD(from_entity=EntityD._from_proto(proto.from_entity),
+                                   relationship=RelationshipD._from_proto(proto.relationship),
+                                   to_entity=EntityD._from_proto(proto.to_entity))
+
+    def to_proto(self) -> entity_pb2.EntityRelationship:
+        return entity_pb2.EntityRelationship(from_entity=self.from_entity.to_proto(),
+                                             relationship=self.relationship.to_proto(),
+                                             to_entity=self.to_entity.to_proto())
+
+    @property
+    def neo4j_create_cmds(self):
+        return [
+            self.from_entity.neo4j_create_cmd,
+            self.to_entity.neo4j_create_cmd,
+            self.relationship.neo4j_create_cmd
+        ]
+
+    @property
+    def neo4j_create_args(self) -> list[Neo4jDict]:
+        relationship_args = {
+            **self.relationship.neo4j_create_args,
+            'from_name': self.from_entity.entity_name,
+            'to_name': self.to_entity.entity_name,
+        }
+
+        return [
+            self.from_entity.neo4j_create_args, self.to_entity.neo4j_create_args, relationship_args
+        ]
+
+
+@dataclasses.dataclass(frozen=True)
+class EntityKnowledgeGraphD(DomainProtocol[entity_pb2.EntityKnowledgeGraph]):
+    entity_relationships: list[EntityRelationshipD]
+
+    @property
+    def id(self) -> str:
+        return hashlib.sha256(self.to_proto().SerializeToString()).hexdigest()
+
+    @classmethod
+    def _from_proto(cls, proto: entity_pb2.EntityKnowledgeGraph) -> EntityKnowledgeGraphD:
+        return EntityKnowledgeGraphD(entity_relationships=[
+            EntityRelationshipD._from_proto(entity_relationship)
+            for entity_relationship in proto.entity_relationships
+        ])
+
+    def to_proto(self) -> entity_pb2.EntityKnowledgeGraph:
+        return entity_pb2.EntityKnowledgeGraph(entity_relationships=[
+            entity_relationship.to_proto() for entity_relationship in self.entity_relationships
+        ])

domain/entity_d_test.py

@@ -0,0 +1,96 @@
+import logging
+import unittest
+
+import proto.entity_pb2 as entity_pb2
+
+import domain.entity_d as entity_d
+
+
+class EntityDTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.entity_d = entity_d.EntityD(entity_id='demoid', entity_name='demo entity')
+
+    def test_proto_roundtrip(self):
+        proto = self.entity_d.to_proto()
+        domain = entity_d.EntityD.from_proto(proto)
+        self.assertEqual(self.entity_d.to_proto(), domain.to_proto())
+
+
+class RelationshipDTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.relationship_d = entity_d.RelationshipD(
+            relationship_id="demoid",
+            start_date="2024-06-01",
+            end_date="2024-06-02",
+            source_text="source text",
+            predicted_movement=entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL)
+
+    def test_proto_roundtrip(self):
+        proto = self.relationship_d.to_proto()
+        domain = entity_d.RelationshipD.from_proto(proto)
+        self.assertEqual(self.relationship_d.to_proto(), domain.to_proto())
+
+    def test_end_date_after_start_date(self):
+        with self.assertRaises(ValueError):
+            _ = entity_d.RelationshipD(
+                relationship_id="demoid",
+                start_date="2024-06-01",
+                end_date="2024-05-02",
+                source_text="source text",
+                predicted_movement=entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL)
+
+
+class EntityRelationshipDTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.from_entity_d = entity_d.EntityD(entity_id='from_id', entity_name='from entity')
+        cls.relationship_d = entity_d.RelationshipD(
+            relationship_id='relationship_id',
+            start_date='2024-06-01',
+            end_date='2024-06-02',
+            source_text='source text',
+            predicted_movement=entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL)
+        cls.to_entity_d = entity_d.EntityD(entity_id='to_id', entity_name='to entity')
+        cls.entity_relationship_d = entity_d.EntityRelationshipD(from_entity=cls.from_entity_d,
+                                                                 relationship=cls.relationship_d,
+                                                                 to_entity=cls.to_entity_d)
+
+    def test_proto_roundtrip(self):
+        proto = self.entity_relationship_d.to_proto()
+        domain = entity_d.EntityRelationshipD.from_proto(proto)
+        self.assertEqual(self.entity_relationship_d.to_proto(), domain.to_proto())
+
+
+class EntityKnowledgeGraphDTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.from_entity_d = entity_d.EntityD(entity_id='from_id', entity_name='from entity')
+        cls.relationship_d = entity_d.RelationshipD(
+            relationship_id='relationship_id',
+            start_date='2024-06-01',
+            end_date='2024-06-02',
+            source_text='source text',
+            predicted_movement=entity_pb2.PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL)
+        cls.to_entity_d = entity_d.EntityD(entity_id='to_id', entity_name='to entity')
+        cls.entity_relationship_d = entity_d.EntityRelationshipD(from_entity=cls.from_entity_d,
+                                                                 relationship=cls.relationship_d,
+                                                                 to_entity=cls.to_entity_d)
+
+        cls.entity_knowledge_graph_d = entity_d.EntityKnowledgeGraphD(
+            entity_relationships=[cls.entity_relationship_d for _ in range(2)])
+
+    def test_proto_roundtrip(self):
+        proto = self.entity_knowledge_graph_d.to_proto()
+        domain = entity_d.EntityKnowledgeGraphD.from_proto(proto)
+        self.assertEqual(self.entity_knowledge_graph_d.to_proto(), domain.to_proto())
+
+
+if __name__ == '__main__':
+    logging.getLogger().setLevel(logging.INFO)
+    unittest.main()

extraction_pipeline/base_stage.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import dataclasses
+import logging
+from typing import Iterable, Protocol, Set, TypeVar
+
+from domain.domain_protocol import DomainProtocol
+
+
+class StageError(Exception):
+    pass
+
+
+StageInputType = TypeVar("StageInputType", bound=DomainProtocol, contravariant=True)
+StageOutputType = TypeVar("StageOutputType", bound=DomainProtocol, covariant=True)
+
+
+@dataclasses.dataclass(frozen=True)
+class BaseStage(Protocol[StageInputType, StageOutputType]):
+
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    def _process_element(self, element: StageInputType) -> Iterable[StageOutputType]:
+        ...
+
+    def process_element(self, element: StageInputType) -> Iterable[StageOutputType]:
+        try:
+            logging.info(f"{self.name} processing element {element}")
+            yield from self._process_element(element)
+        except Exception as e:
+            logging.error(e)
+            raise StageError(f"Error processing element {element}") from e
+
+
+class TransformError(Exception):
+    pass
+
+
+TransformInputType = TypeVar("TransformInputType", bound=DomainProtocol, contravariant=True)
+TransformOutputType = TypeVar("TransformOutputType", bound=DomainProtocol, covariant=True)
+
+
+@dataclasses.dataclass(frozen=True)
+class BaseTransform(Protocol[TransformInputType, TransformOutputType]):
+
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    def _process_collection(
+            self, collection: Iterable[TransformInputType]) -> Iterable[TransformOutputType]:
+        ...
+
+    def process_collection(
+            self, collection: Iterable[TransformInputType]) -> Iterable[TransformOutputType]:
+        try:
+            logging.info(f"{self.name} processing collection")
+            yield from self._process_collection(collection)
+        except Exception as e:
+            logging.error(e)
+            raise TransformError(f"Error processing collection") from e

extraction_pipeline/document_metadata_extractor/document_metadata_extractor.py

@@ -0,0 +1,9 @@
+import dataclasses
+
+from extraction_pipeline.base_stage import BaseStage
+from domain.chunk_d import DocumentD
+
+
+@dataclasses.dataclass(frozen=True)
+class DocumentMetadataExtractor(BaseStage[DocumentD, DocumentD]):
+    ...

extraction_pipeline/document_metadata_extractor/openai_document_metadata_extractor.py

@@ -0,0 +1,73 @@
+import dataclasses
+import json
+from typing import ClassVar, Dict, Iterable, Iterator, List, Optional, Tuple, Union
+
+import pymupdf
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
+
+from domain.chunk_d import DocumentD
+from extraction_pipeline.document_metadata_extractor.document_metadata_extractor import (
+    DocumentMetadataExtractor,)
+from extraction_pipeline.document_metadata_extractor.prompts import (
+    DOCUMENT_METADATA_PROMPT,)
+from llm_handler.llm_interface import LLMInterface
+from llm_handler.openai_handler import ChatModelVersion, OpenAIHandler
+from utils.dates import parse_date
+
+
+class CreationDateError(Exception):
+    pass
+
+
+class AuthorsError(Exception):
+    pass
+
+
+class OpenAIDocumentMetadataExtractor(DocumentMetadataExtractor):
+
+    _handler: LLMInterface
+    _MODEL_VERSION: ChatModelVersion = ChatModelVersion.GPT_4_O
+    _AUTHORS_KEY: ClassVar[str] = "authors"
+    _PUBLISH_DATE_KEY: ClassVar[str] = "publish_date"
+    _TEMPARATURE: ClassVar[float] = 0.2
+
+    def __init__(self,
+                 openai_handler: Optional[LLMInterface] = None,
+                 model_version: Optional[ChatModelVersion] = None):
+        self._handler = openai_handler or OpenAIHandler()
+        self._model_version = model_version or self._MODEL_VERSION
+
+    def _validate_text(self, completion_text: Dict[str, Union[str, List[str]]]):
+        if not completion_text.get(self._AUTHORS_KEY):
+            raise AuthorsError("No authors found.")
+        if not completion_text.get(self._PUBLISH_DATE_KEY):
+            raise CreationDateError("No creation date found.")
+
+        publish_date_str: str = str(completion_text.get(self._PUBLISH_DATE_KEY, ""))
+        try:
+            parse_date(publish_date_str)
+        except ValueError as e:
+            raise CreationDateError(
+                f"Failed to parse publish date '{publish_date_str}': {e}") from e
+
+    def _process_element(self, element: DocumentD) -> Iterable[DocumentD]:
+        pdf_document_pages: Iterator = pymupdf.open(element.file_path).pages()
+        first_page_text: str = next(pdf_document_pages).get_text()
+        messages: List[ChatCompletionMessageParam] = [{
+            "role": "system", "content": DOCUMENT_METADATA_PROMPT
+        },
+                                                      {
+                                                          "role": "user",
+                                                          "content": f"Input:\n{first_page_text}"
+                                                      }]
+        completion_text_raw = self._handler.get_chat_completion(
+            messages=messages,
+            model=self._model_version,
+            temperature=self._TEMPARATURE,
+            response_format={"type": "json_object"})
+        completion_text: Dict[str, Union[str, List[str]]] = dict(json.loads(completion_text_raw))
+        self._validate_text(completion_text)
+        authors: str = ", ".join(completion_text.get(self._AUTHORS_KEY, []))
+        publish_date: str = str(completion_text.get(self._PUBLISH_DATE_KEY, ""))
+
+        yield dataclasses.replace(element, authors=authors, publish_date=publish_date)

extraction_pipeline/document_metadata_extractor/openai_document_metadata_extractor_test.py

@@ -0,0 +1,74 @@
+import logging
+import unittest
+import uuid
+from typing import List
+
+from domain.chunk_d import DocumentD
+from extraction_pipeline.document_metadata_extractor.openai_document_metadata_extractor import OpenAIDocumentMetadataExtractor, AuthorsError, CreationDateError
+from llm_handler.mock_llm_handler import MockLLMHandler
+
+DOCUMENT_METADATA_EXTRACTION_RESPONSE = '''
+{
+    "authors": ["BofA Global Research", "Michael Hartnett", "Elyas Galou", "Anya Shelekhin", "Myung-Jee Jung"],
+    "publish_date": "2023-04-13"
+}
+'''
+
+
+class TestOpenAIDocumentMetadataExtractor(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.test_pdf_path = "extraction_pipeline/test_data/test.pdf"
+        cls.start_document_d = DocumentD(file_path=cls.test_pdf_path, authors="", publish_date="")
+        cls.final_document_d = DocumentD(
+            file_path=cls.test_pdf_path,
+            authors=
+            "BofA Global Research, Michael Hartnett, Elyas Galou, Anya Shelekhin, Myung-Jee Jung",
+            publish_date="2023-04-13")
+        cls.openai_publish_details_extractor = OpenAIDocumentMetadataExtractor()
+
+    def test__validate_text_missing_publishers(self):
+        missing_publishers_text = {"publish_date": "2023-12-13"}
+        with self.assertRaises(AuthorsError):
+            self.openai_publish_details_extractor._validate_text(
+                missing_publishers_text)  # type: ignore
+
+    def test__validate_text_invalid_date(self):
+        invalid_date_text = {
+            "authors": [
+                "BofA Global Research",
+                "Michael Hartnett",
+                "Elyas Galou",
+                "Anya Shelekhin",
+                "Myung-Jee Jung"
+            ],
+            "publish_date": "2-13"
+        }
+        with self.assertRaises(CreationDateError):
+            self.openai_publish_details_extractor._validate_text(invalid_date_text)  # type: ignore
+
+    def test__validate_text_valid(self):
+        valid_text = {
+            "authors": [
+                "BofA Global Research",
+                "Michael Hartnett",
+                "Elyas Galou",
+                "Anya Shelekhin",
+                "Myung-Jee Jung"
+            ],
+            "publish_date": "2023-04-13"
+        }
+        self.openai_publish_details_extractor._validate_text(valid_text)  # type: ignore
+
+    def test__process_element(self):
+        handler = MockLLMHandler(chat_completion=[DOCUMENT_METADATA_EXTRACTION_RESPONSE])
+        openai_publish_details_extractor = OpenAIDocumentMetadataExtractor(handler)
+        pdf_document_d = self.start_document_d
+        output = list(openai_publish_details_extractor._process_element(pdf_document_d))
+        self.assertEqual(output[0], self.final_document_d)
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    unittest.main()

extraction_pipeline/document_metadata_extractor/prompts.py

@@ -0,0 +1,22 @@
+DOCUMENT_METADATA_PROMPT = '''
+You are a financial reports expert, tasked with extracting the authors and publishing date
+from the first page of a financial report document. Given the first page of text from a financial report,
+you must extract and return the list of authors (which may be only one) and the date of
+publication in a specific string format of 'YYYY-MM-DD' in the form of a JSON object. The JSON object
+you return should have the list of authors as the value to the "authors" key. 
+The specifically formatted datetime string should be the value to the "publish_date" key. 
+Note, I will be parsing out the datetime string using the python function `datetime.strptime(your_output['publish_date'], '%Y-%m-%d')`. 
+Do not include anything besides the JSON object in your response. 
+I will be parsing your response using the python code `json.loads(your_output)`.
+
+Example input:
+"Trading ideas and investment strategies discussed herein may give rise to significant risk and are \nnot suitable for all investors. Investors should have experience in relevant markets and the financial \nresources to absorb any losses arising from applying these ideas or strategies. \n>> Employed by a non-US affiliate of BofAS and is not registered/qualified as a research analyst \nunder the FINRA rules. \nRefer to "Other Important Disclosures" for information on certain BofA Securities entities that take \nresponsibility for the information herein in particular jurisdictions. \nBofA Securities does and seeks to do business with issuers covered in its research \nreports. As a result, investors should be aware that the firm may have a conflict of \ninterest that could affect the objectivity of this report. Investors should consider this \nreport as only a single factor in making their investment decision. \nRefer to important disclosures on page 11 to 13.  \n12544294 \n \nThe Flow Show \n \nSugar Coated Iceberg \n \n \nScores on the Doors: crypto 65.3%, gold 10.1%, stocks 7.7%, HY bonds 4.3%, IG bonds \n4.3%, govt bonds 3.4%, oil 3.7%, cash 1.2%, commodities -0.1%, US dollar -2.0% YTD. \nThe Biggest Picture: everyone's new favorite theme…US dollar debasement; US$ -11% \nsince Sept, gold >$2k, bitcoin >$30k; right secular theme (deficits, debt, geopolitics), \nUS$ in 4th bear market of past 50 years, bullish gold, oil, Euro, international stocks; but \npessimism so high right now, and 200bps Fed cuts priced in from June to election, look \nfor US$ trading bounce after Fed ends hiking cycle May 3rd (Chart 2). \nTale of the Tape: inflation slowing, Fed hiking cycle over, recession expectations \nuniversal, yet UST 30-year can't break below 3.6% (200-dma) because we've already \ntraded 8% to 4% CPI, labor market yet to crack, US govt deficit growing too quickly.  \nThe Price is Right: next 6 months it's "recession vs Fed cuts"; best tells who's \nwinning…HY bonds, homebuilders, semiconductors…HYG <73, XHB <70, SOX <2900 \nrecessionary, and if levels hold…it's a no/soft landing.  \nChart 2: US dollar has begun 4th bear market of past 50 years \nUS dollar index (DXY) \n \nSource: BofA Global Investment Strategy, Bloomberg. \nBofA GLOBAL RESEARCH \nMore on page 2… \n \n \n \n \n \n \nGold ↑\nYen & Deutsche \nMark ↑\nBRICS & Euro ↑\nGold, \nEuro ↑\n70\n80\n90\n100\n110\n120\n130\n140\n150\n160\n1967\n1973\n1979\n1985\n1991\n1997\n2003\n2009\n2015\n2021\nUS dollar index (DXY)\n13 April 2023 Corrected \n \n \n \n \nInvestment Strategy \nGlobal \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \nMichael Hartnett \nInvestment Strategist \nBofAS \n+1 646 855 1508 \nmichael.hartnett@bofa.com \n \nElyas Galou >> \nInvestment Strategist \nBofASE (France) \n+33 1 8770 0087 \nelyas.galou@bofa.com \n \nAnya Shelekhin \nInvestment Strategist \nBofAS \n+1 646 855 3753 \nanya.shelekhin@bofa.com \n \nMyung-Jee Jung \nInvestment Strategist \nBofAS \n+1 646 855 0389 \nmyung-jee.jung@bofa.com \n \n \n \n \n \n \n \n \n \nChart 1: BofA Bull & Bear Indicator \nStays at 2.3 \n \nSource: BofA Global Investment Strategy \nThe indicator identified above as the BofA Bull & Bear \nIndicator is intended to be an indicative metric only and \nmay not be used for reference purposes or as a measure \nof performance for any financial instrument or contract, \nor otherwise relied upon by third parties for any other \npurpose, without the prior written consent of BofA \nGlobal Research. This indicator was not created to act as \na benchmark. \nBofA GLOBAL RESEARCH \n \n \nExtreme \nBearish\nExtreme \nBullish\n4\n6\n10\n0\n8\n2\nBuy\nSell\n2.3\nAccessible version \n \n \nTimestamp: 13 April 2023 08:19PM EDT\nW \n"
+
+Example output:
+"
+{
+     "authors": ["BofA Global Research", "Michael Hartnett", "Elyas Galou", "Anya Shelekhin", "Myung-Jee Jung"],
+     "publish_date": "2023-04-13"
+}
+"
+'''

extraction_pipeline/pdf_process_stage.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+import os
+from typing import Iterable
+import pymupdf
+import re
+
+from domain.domain_protocol import DomainProtocol
+from domain.chunk_d import DocumentD, ChunkD
+import proto.chunk_pb2 as chunk_pb2
+from extraction_pipeline.base_stage import BaseStage, BaseTransform
+from storage.domain_dao import CacheDomainDAO
+
+
+class IDTagger():
+
+    _uuid_to_tag: CacheDomainDAO[ChunkD]
+
+    def __init__(self, cache_save_path: str):
+        self._uuid_to_tag = CacheDomainDAO(f"{cache_save_path}.json", ChunkD)
+
+    def __call__(self, element: ChunkD) -> ChunkD:
+        self._uuid_to_tag.insert([element])
+        return element
+
+
+class PdfPageChunkerStage(BaseStage[DocumentD, ChunkD]):
+
+    def __init__(self):
+        self._pdf_cache_dao = CacheDomainDAO("pdf_chunk_id_map.json", DocumentD)
+        self.id_tagger = CacheDomainDAO(f"{str.lower((self.__class__.__name__))}_chunk_id_map.json",
+                                        ChunkD)
+
+    def _process_element(self, element: DocumentD) -> Iterable[ChunkD]:
+        pdf_document: Iterable = pymupdf.open(element.file_path)  # type: ignore
+        for page_index, pdf_page in enumerate(pdf_document):
+            page_text: str = pdf_page.get_textpage().extractText()  # type: ignore
+            yield self.id_tagger(
+                ChunkD(parent_reference=element,
+                       chunk_text=page_text,
+                       chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PAGE,
+                       chunk_index=page_index,
+                       chunk_id=self._pdf_cache_dao.set(element)))
+
+
+class ParagraphChunkerStage(BaseStage[ChunkD, ChunkD]):
+
+    def __init__(self):
+        self.id_tagger = CacheDomainDAO(f"{str.lower((self.__class__.__name__))}_chunk_id_map.json",
+                                        ChunkD)
+
+    def _process_element(self, element: ChunkD) -> Iterable[ChunkD]:
+        paragraphs = re.split(r'\n+', element.chunk_text)
+        paragraphs = [para.strip() for para in paragraphs if para.strip()]
+        for chunk_index, paragraph in enumerate(paragraphs):
+            yield self.id_tagger(
+                ChunkD(parent_reference=element.chunk_id,
+                       chunk_text=paragraph,
+                       chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PARAGRAPH,
+                       chunk_index=chunk_index))
+
+
+class SentenceChunkerStage(BaseStage[ChunkD, ChunkD]):
+
+    def __init__(self):
+        self.id_tagger = CacheDomainDAO(f"{str.lower((self.__class__.__name__))}_chunk_id_map.json",
+                                        ChunkD)
+
+    def _process_element(self, element: ChunkD) -> Iterable[ChunkD]:
+        sentence_endings = r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?|!)(\s|$)'
+        sentences = re.split(sentence_endings, element.chunk_text)
+        sentences = [sentence.strip() for sentence in sentences if sentence.strip()]
+        for chunk_index, sentence in enumerate(sentences):
+            yield self.id_tagger(
+                ChunkD(parent_reference=element.chunk_id,
+                       chunk_text=sentence,
+                       chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_SENTENCE,
+                       chunk_index=chunk_index))
+
+
+class PdfToSentencesTransform(BaseTransform[DocumentD, ChunkD]):
+
+    _pdf_page_chunker: PdfPageChunkerStage
+    _paragraph_chunker: ParagraphChunkerStage
+    _sentence_chunker: SentenceChunkerStage
+
+    def __init__(self):
+        self._pdf_page_chunker = PdfPageChunkerStage()
+        self._paragraph_chunker = ParagraphChunkerStage()
+        self._sentence_chunker = SentenceChunkerStage()
+
+    def _process_collection(self, collection: Iterable[DocumentD]) -> Iterable[ChunkD]:
+        for pdf_document in collection:
+            pdf_pages = self._pdf_page_chunker.process_element(pdf_document)
+            for pdf_page in pdf_pages:
+                paragraphs = self._paragraph_chunker.process_element(pdf_page)
+                for paragraph in paragraphs:
+                    sentences = self._sentence_chunker.process_element(paragraph)
+                    yield from sentences
+
+
+class PdfToParagraphTransform(BaseTransform[DocumentD, ChunkD]):
+
+    _pdf_page_chunker: PdfPageChunkerStage
+    _paragraph_chunker: ParagraphChunkerStage
+
+    def __init__(self):
+        self._pdf_page_chunker = PdfPageChunkerStage()
+        self._paragraph_chunker = ParagraphChunkerStage()
+
+    def _process_collection(self, collection: Iterable[DocumentD]) -> Iterable[ChunkD]:
+        for pdf_document in collection:
+            pdf_pages = self._pdf_page_chunker.process_element(pdf_document)
+            for pdf_page in pdf_pages:
+                paragraphs = self._paragraph_chunker.process_element(pdf_page)
+                for paragraph in paragraphs:
+                    yield paragraph
+
+
+class PdfToPageTransform(BaseTransform[DocumentD, ChunkD]):
+
+    _pdf_page_chunker: PdfPageChunkerStage
+
+    def __init__(self):
+        self._pdf_page_chunker = PdfPageChunkerStage()
+        self._paragraph_chunker = ParagraphChunkerStage()
+
+    def _process_collection(self, collection: Iterable[DocumentD]) -> Iterable[ChunkD]:
+        for pdf_document in collection:
+            pdf_pages = self._pdf_page_chunker.process_element(pdf_document)
+            for pdf_page in pdf_pages:
+                yield pdf_page

extraction_pipeline/pdf_process_stage_test.py

@@ -0,0 +1,103 @@
+import unittest
+import logging
+import uuid
+
+from domain.chunk_d import DocumentD, ChunkD
+import proto.chunk_pb2 as chunk_pb2
+from extraction_pipeline.pdf_process_stage import PdfPageChunkerStage, ParagraphChunkerStage, SentenceChunkerStage, PdfToSentencesTransform
+
+
+def compare_chunk_without_id(chunk1: ChunkD, chunk2: ChunkD):
+    assert chunk1.parent_reference == chunk2.parent_reference
+    assert chunk1.chunk_text == chunk2.chunk_text
+    assert chunk1.chunk_type == chunk2.chunk_type
+    assert chunk1.chunk_index == chunk2.chunk_index
+
+
+class PdfPageChunkerStageTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.test_pdf_path = "extraction_pipeline/test_data/test.pdf"
+        cls.expected_pdf_text = """This research presents a thorough examination of Large Language Models (LLMs) in the
+domain of code re-engineering.
+Our focus on the methodical development and utilization of robust pipelines highlights both the
+potential and the current challenges of employing LLMs for these complex tasks.
+"""
+        cls.document_d = DocumentD(file_path="extraction_pipeline/test_data/test.pdf",
+                                   authors="BofA John Doe Herb Johnson Taylor Mason",
+                                   publish_date="2021-01-01")
+
+    def test__process_element(self):
+        chunks = list(PdfPageChunkerStage()._process_element(self.document_d))
+        self.assertEqual(len(chunks), 1)
+        expected_chunk = ChunkD(parent_reference=self.document_d,
+                                chunk_text=self.expected_pdf_text,
+                                chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PAGE,
+                                chunk_index=0)
+        compare_chunk_without_id(chunks[0], expected_chunk)
+
+
+class ParagraphChunkerStageTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.document_d = DocumentD(file_path="extraction_pipeline/test_data/test.pdf",
+                                   authors="BofA John Doe Herb Johnson Taylor Mason",
+                                   publish_date="2021-01-01")
+
+    def test__process_element(self):
+        chunk = ChunkD(parent_reference=self.document_d,
+                       chunk_text="paragraph1\nparagraph2\n\nparagraph3",
+                       chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PAGE,
+                       chunk_index=0)
+        chunks = list(ParagraphChunkerStage()._process_element(chunk))
+        expected_chunks = [
+            ChunkD(parent_reference=chunk.chunk_id,
+                   chunk_text="paragraph1",
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PARAGRAPH,
+                   chunk_index=0),
+            ChunkD(parent_reference=chunk.chunk_id,
+                   chunk_text="paragraph2",
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PARAGRAPH,
+                   chunk_index=1),
+            ChunkD(parent_reference=chunk.chunk_id,
+                   chunk_text="paragraph3",
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PARAGRAPH,
+                   chunk_index=2)
+        ]
+        self.assertEqual(len(chunks), len(expected_chunks))
+        for chunk, expected_chunk in zip(chunks, expected_chunks):
+            compare_chunk_without_id(chunk, expected_chunk)
+
+
+class SentenceChunkerStageTest(unittest.TestCase):
+
+    def test__process_element(self):
+        chunk = ChunkD(parent_reference=uuid.uuid4(),
+                       chunk_text="sentence1. sentence2! sentence3?",
+                       chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PARAGRAPH,
+                       chunk_index=0)
+        chunks = list(SentenceChunkerStage()._process_element(chunk))
+        expected_chunks = [
+            ChunkD(parent_reference=chunk.chunk_id,
+                   chunk_text="sentence1.",
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_SENTENCE,
+                   chunk_index=0),
+            ChunkD(parent_reference=chunk.chunk_id,
+                   chunk_text="sentence2!",
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_SENTENCE,
+                   chunk_index=1),
+            ChunkD(parent_reference=chunk.chunk_id,
+                   chunk_text="sentence3?",
+                   chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_SENTENCE,
+                   chunk_index=2)
+        ]
+        self.assertEqual(len(chunks), len(expected_chunks))
+        for chunk, expected_chunk in zip(chunks, expected_chunks):
+            compare_chunk_without_id(chunk, expected_chunk)
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    unittest.main()

extraction_pipeline/pdf_to_knowledge_graph_transform.py

@@ -0,0 +1,153 @@
+import argparse
+import logging
+import os
+import sys
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Iterable, Optional
+
+from domain.chunk_d import ChunkD, DocumentD
+from domain.entity_d import EntityKnowledgeGraphD
+from extraction_pipeline.base_stage import BaseTransform
+from extraction_pipeline.document_metadata_extractor.openai_document_metadata_extractor import (
+    OpenAIDocumentMetadataExtractor,)
+from extraction_pipeline.pdf_process_stage import (
+    PdfToPageTransform,
+    PdfToParagraphTransform,
+    PdfToSentencesTransform,
+)
+from extraction_pipeline.relationship_extractor.entity_relationship_extractor import (
+    RelationshipExtractor,)
+from extraction_pipeline.relationship_extractor.openai_relationship_extractor import (
+    OpenAIRelationshipExtractor,)
+from llm_handler.openai_handler import OpenAIHandler
+from storage.domain_dao import InMemDomainDAO
+from storage.neo4j_dao import Neo4jDomainDAO
+
+
+class PdfToKnowledgeGraphTransform(BaseTransform[DocumentD, EntityKnowledgeGraphD]):
+    _metadata_extractor: OpenAIDocumentMetadataExtractor
+    _pdf_chunker: BaseTransform[DocumentD, ChunkD]
+    _relationship_extractor: RelationshipExtractor
+
+    def __init__(self, pdf_chunker: BaseTransform[DocumentD, ChunkD]):
+        openai_handler = OpenAIHandler()
+        self._metadata_extractor = OpenAIDocumentMetadataExtractor(openai_handler=openai_handler)
+        self._pdf_chunker = pdf_chunker
+        self._relationship_extractor = OpenAIRelationshipExtractor(openai_handler=openai_handler)
+
+    def _process_collection(self,
+                            collection: Iterable[DocumentD]) -> Iterable[EntityKnowledgeGraphD]:
+        # produce 1 EntityKnowledgeGraphD per DocumentD
+        for pdf_document in collection:
+            # metadata extractor only yields 1 filled in DocumentD for each input DocumentD
+            document = next(iter(self._metadata_extractor.process_element(pdf_document)))
+            entity_relationships = []
+            for pdf_chunk in self._pdf_chunker.process_collection([document]):
+                for relationship in self._relationship_extractor.process_element(pdf_chunk):
+                    entity_relationships.append(relationship)
+            yield EntityKnowledgeGraphD(entity_relationships=entity_relationships)
+
+
+if __name__ == '__main__':
+    ## CLI Arguments for running transform as multi-threaded script
+    parser = argparse.ArgumentParser(description='Extract knowledge graphs from PDF files')
+    parser.add_argument('--pdf_folder',
+                        type=str,
+                        help='Path to folder of PDF files to process',
+                        default='')
+    parser.add_argument('--pdf_file',
+                        type=str,
+                        help='Path to the one PDF file to process',
+                        default='')
+    parser.add_argument('--output_json_file',
+                        type=str,
+                        help='Path for output json file of knowledge graphs',
+                        default='./knowledge_graphs.json')
+    parser.add_argument('--log_folder', type=str, help='Path to log folder', default='log')
+    parser.add_argument('--chunk_to',
+                        type=str,
+                        help='What level to chunk PDF text',
+                        default='page',
+                        choices=['page', 'paragraph', 'sentence'])
+    parser.add_argument('--verbose', help='Enable DEBUG level logs', action='store_true')
+    parser.add_argument('--upload_to_neo4j',
+                        help='Enable uploads to Neo4j database',
+                        action='store_true')
+    args = parser.parse_args()
+
+    ## Setup logging
+    if args.verbose:
+        log_level = logging.DEBUG
+    else:
+        log_level = logging.INFO
+    os.makedirs(args.log_folder, exist_ok=True)
+    script_name = os.path.splitext(os.path.basename(__file__))[0]
+    logging.basicConfig(level=log_level,
+                        format='%(asctime)s - %(levelname)s - %(message)s',
+                        filename=f'{args.log_folder}/{script_name}.log',
+                        filemode='w')
+    logger = logging.getLogger(__name__)
+
+    ## Setup PDF Chunking
+    if args.chunk_to == 'page':
+        pdf_transform = PdfToPageTransform()
+    elif args.chunk_to == 'paragraph':
+        pdf_transform = PdfToParagraphTransform()
+    elif args.chunk_to == 'sentence':
+        pdf_transform = PdfToSentencesTransform()
+    else:
+        logging.error('Invalid chunking level: %s', args.chunk_to)
+        sys.exit(1)
+
+    ## Process PDF Files
+    if not args.pdf_folder and not args.pdf_file:
+        logging.error('No PDF file or folder provided')
+        sys.exit(1)
+    elif args.pdf_folder:
+        pdf_folder = Path(args.pdf_folder)
+        if not pdf_folder.exists():
+            logging.error('PDF folder does not exist: %s', pdf_folder)
+            sys.exit(1)
+        pdf_files = list(pdf_folder.glob('*.pdf'))
+        if len(pdf_files) == 0:
+            logging.warning('No PDF files found in folder: %s', pdf_folder)
+            sys.exit(0)
+        pdfs = [
+            DocumentD(file_path=str(pdf_file), authors='', publish_date='')
+            for pdf_file in pdf_files
+        ]
+    else:
+        pdf_file = Path(args.pdf_file)
+        if not pdf_file.exists():
+            logging.error('PDF file does not exist: %s', pdf_file)
+            sys.exit(1)
+        pdfs = [DocumentD(file_path=str(pdf_file), authors='', publish_date='')]
+
+    pdf_to_kg = PdfToKnowledgeGraphTransform(pdf_transform)
+
+    def process_pdf(pdf: DocumentD) -> tuple[Optional[EntityKnowledgeGraphD], str]:
+        pdf_name = Path(pdf.file_path).name
+        try:
+            # process collection yields 1 KG per pdf but we are only
+            # inputing 1 PDF at a time so we just need the 1st element
+            return list(pdf_to_kg.process_collection([pdf]))[0], pdf_name
+        except Exception as e:
+            logging.error(f"Error processing pdf: {e}")
+            return None, pdf_name
+
+    results: list[EntityKnowledgeGraphD] = []
+    with ThreadPoolExecutor() as executor, Neo4jDomainDAO() as dao:
+        futures = [executor.submit(process_pdf, pdf) for pdf in pdfs]
+
+        for future in as_completed(futures):
+            kg, pdf_name = future.result()
+            if not kg:
+                continue
+            results.append(kg)
+            if args.upload_to_neo4j:
+                dao.insert(kg, pdf_name)
+
+    dao = InMemDomainDAO()
+    dao.insert(results)
+    dao.save_to_file(args.output_json_file)

extraction_pipeline/relationship_extractor/entity_relationship_extractor.py

@@ -0,0 +1,10 @@
+import dataclasses
+
+from extraction_pipeline.base_stage import BaseStage
+from domain.chunk_d import ChunkD
+from domain.entity_d import EntityRelationshipD
+
+
+@dataclasses.dataclass(frozen=True)
+class RelationshipExtractor(BaseStage[ChunkD, EntityRelationshipD]):
+    ...

extraction_pipeline/relationship_extractor/openai_relationship_extractor.py

@@ -0,0 +1,108 @@
+import json
+import logging
+import uuid
+from typing import ClassVar, Dict, Iterator, List, Optional
+
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
+
+from domain.chunk_d import (ChunkD, DocumentD)
+from domain.entity_d import (
+    EntityD,
+    EntityRelationshipD,
+    RelationshipD,
+)
+from extraction_pipeline.relationship_extractor.entity_relationship_extractor import (
+    RelationshipExtractor,)
+from extraction_pipeline.relationship_extractor.prompts import (
+    EXTRACT_RELATIONSHIPS_PROMPT,
+    NER_TAGGING_PROMPT,
+)
+from llm_handler.llm_interface import LLMInterface
+from llm_handler.openai_handler import ChatModelVersion, OpenAIHandler
+
+
+class OpenAIRelationshipExtractor(RelationshipExtractor):
+
+    _handler: LLMInterface
+    _MODEL_VERSION: ChatModelVersion = ChatModelVersion.GPT_4_O
+    _RELATIONSHIP_KEY: ClassVar[str] = "relationships"
+    _ENTITIES_KEY: ClassVar[str] = "entities"
+    _RELATIONSHIPS_TYPES: ClassVar[List[str]] = ["PREDICTION"]
+    _TEMPARATURE: ClassVar[float] = 0.2
+
+    def __init__(self,
+                 openai_handler: Optional[LLMInterface] = None,
+                 model_version: Optional[ChatModelVersion] = None):
+        self._handler = openai_handler or OpenAIHandler()
+        self._model_version = model_version or self._MODEL_VERSION
+
+    def _extract_entity_names(self, chunk_text: str) -> List[Dict[str, str]]:
+
+        messages: List[ChatCompletionMessageParam] = [{
+            "role": "system", "content": NER_TAGGING_PROMPT
+        },
+                                                      {
+                                                          "role": "user",
+                                                          "content": f"Input:\n{chunk_text}"
+                                                      }]
+        completion_text = self._handler.get_chat_completion(messages=messages,
+                                                            model=self._model_version,
+                                                            temperature=self._TEMPARATURE,
+                                                            response_format={"type": "json_object"})
+        logging.info(f"entity extraction results: {completion_text}")
+        return dict(json.loads(completion_text)).get(self._ENTITIES_KEY, [])
+
+    def _extract_relationships(self, chunk: ChunkD,
+                               entity_nodes: List[Dict[str, str]]) -> Iterator[EntityRelationshipD]:
+        if isinstance(chunk.parent_reference, DocumentD):
+            analyst_names: str = chunk.parent_reference.authors
+            document_datetime: str = chunk.parent_reference.publish_date
+        else:
+            raise NotImplementedError("Parent reference is not a DocumentD")
+        messages: List[ChatCompletionMessageParam] = [
+            {
+                "role": "system", "content": EXTRACT_RELATIONSHIPS_PROMPT
+            },
+            {
+                "role":
+                    "user",
+                "content":
+                    f"Analyst: {analyst_names} \n Date: {document_datetime} \n Text Chunk: {chunk.chunk_text} \n {str(entity_nodes)}"
+            }
+        ]
+
+        completion_text = self._handler.get_chat_completion(messages=messages,
+                                                            model=self._model_version,
+                                                            temperature=self._TEMPARATURE,
+                                                            response_format={"type": "json_object"})
+        logging.info(f"relationship results: {completion_text}")
+        completion_text = dict(json.loads(completion_text))
+        relationships: List[Dict[str, Dict[str,
+                                           str]]] = completion_text.get(self._RELATIONSHIP_KEY, [])
+        for extracted_relationship in relationships:
+            key: str = list(extracted_relationship.keys())[0]
+            if key in self._RELATIONSHIPS_TYPES:
+                relationship_kw_attr: Dict[str, str] = extracted_relationship[key]
+                relationship_d = RelationshipD(
+                    relationship_id=str(uuid.uuid4()),
+                    start_date=relationship_kw_attr.get("start_date", ""),
+                    end_date=relationship_kw_attr.get("end_date", ""),
+                    source_text=chunk.chunk_text,
+                    predicted_movement=RelationshipD.from_string(
+                        relationship_kw_attr.get("predicted_movement", "")))
+            else:
+                raise ValueError(f"No valid relationships in {extracted_relationship}")
+
+            entity_l_d = EntityD(entity_id=str(uuid.uuid4()),
+                                 entity_name=relationship_kw_attr.get("from_entity", ""))
+
+            entity_r_d = EntityD(entity_id=str(uuid.uuid4()),
+                                 entity_name=relationship_kw_attr.get("to_entity", ""))
+
+            yield EntityRelationshipD(relationship=relationship_d,
+                                      from_entity=entity_l_d,
+                                      to_entity=entity_r_d)
+
+    def _process_element(self, element: ChunkD) -> Iterator[EntityRelationshipD]:
+        entities_text: List[Dict[str, str]] = self._extract_entity_names(element.chunk_text)
+        yield from self._extract_relationships(element, entities_text)

extraction_pipeline/relationship_extractor/openai_relationship_extractor_test.py

@@ -0,0 +1,87 @@
+import logging
+import unittest
+import uuid
+
+import proto.chunk_pb2 as chunk_pb2
+
+from domain.chunk_d import ChunkD, DocumentD
+from domain.entity_d import RelationshipD
+from extraction_pipeline.relationship_extractor.openai_relationship_extractor import (
+    OpenAIRelationshipExtractor,)
+from llm_handler.mock_llm_handler import MockLLMHandler
+
+ENTITY_EXTRACTION_RESPONSE = '''
+{
+"entities": [
+    {"name": "entity 1", "type": "Financial Securities"},
+    {"name": "entity 2", "type": "Financial Securities"}
+]
+}
+'''
+
+RELATIONSHIP_EXTRACTION_RESPONSE = '''
+{
+    "relationships": [
+        {
+            "PREDICTION": {
+                "from_entity": "John Doe",
+                "to_entity": "entity 1",
+                "start_date": "2024-06-12",
+                "end_date": "2024-11-05",
+                "predicted_movement": "PREDICTED_MOVEMENT_DECREASE",
+                "description": "200bps Fed cuts priced in from June to election, look for US$ trading bounce after Fed ends hiking cycle May 3rd."
+            }
+        },
+        {
+            "PREDICTION": {
+                "from_entity": "John Doe",
+                "to_entity": "entity 2",
+                "start_date": "2024-06-12",
+                "end_date": "2024-11-05",
+                "predicted_movement": "PREDICTED_MOVEMENT_DECREASE",
+                "description": "200bps Fed cuts priced in from June to election, look for US$ trading bounce after Fed ends hiking cycle May 3rd."
+            }
+        }
+    ]
+}
+'''
+
+
+class TestOpenAIRelationshipExtractor(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.document_d = DocumentD(file_path="Morgan_Stanley_Research_2024.pdf",
+                                    authors='John Doe',
+                                    publish_date='2024-06-12')
+        self.chunkd = ChunkD(chunk_text='entity 1 relates to entity 2.',
+                             chunk_type=chunk_pb2.ChunkType.CHUNK_TYPE_PAGE,
+                             chunk_index=1,
+                             parent_reference=self.document_d)
+
+    def test__process_elements(self):
+        chat_responses = [ENTITY_EXTRACTION_RESPONSE, RELATIONSHIP_EXTRACTION_RESPONSE]
+        handler = MockLLMHandler(chat_completion=chat_responses)
+        openai_relationship_extractor = OpenAIRelationshipExtractor(handler)
+        output = list(openai_relationship_extractor._process_element(self.chunkd))
+        extracted_relationship1 = output[0]
+        extracted_relationship2 = output[1]
+        self.assertEqual(len(output), 2)
+        self.assertEqual(extracted_relationship1.from_entity.entity_name, 'John Doe')
+        self.assertEqual(extracted_relationship1.to_entity.entity_name, 'entity 1')
+        self.assertEqual(extracted_relationship1.relationship.start_date, '2024-06-12')
+        self.assertEqual(extracted_relationship1.relationship.end_date, '2024-11-05')
+        self.assertEqual(extracted_relationship1.relationship.predicted_movement, 3)
+        self.assertEqual(extracted_relationship1.relationship.source_text,
+                         'entity 1 relates to entity 2.')
+        self.assertEqual(extracted_relationship2.from_entity.entity_name, 'John Doe')
+        self.assertEqual(extracted_relationship2.to_entity.entity_name, 'entity 2')
+        self.assertEqual(extracted_relationship2.relationship.start_date, '2024-06-12')
+        self.assertEqual(extracted_relationship2.relationship.end_date, '2024-11-05')
+        self.assertEqual(extracted_relationship2.relationship.predicted_movement, 3)
+        self.assertEqual(extracted_relationship2.relationship.source_text,
+                         'entity 1 relates to entity 2.')
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    unittest.main()

extraction_pipeline/relationship_extractor/prompts.py

@@ -0,0 +1,481 @@
+NER_TAGGING_PROMPT = '''
+You are an expert in finance who will be performing Named Entity Recognition (NER) 
+with the high-level objective of analyzing financial documents. The output should 
+be formatted as a JSON object, where the JSON object has a key "entities" that 
+contains a list of the extracted entities. If no entities are found, return an 
+empty JSON object.
+
+- Finacial Markets (Grouping of financial securities)
+- Finacial Securities
+- Finacial Indicators
+
+Now, given a text chunk, extract the relevant entities and format the output as a JSON object.
+
+Examples of entities:
+
+- Financial Markets: "US Stock Market", "US Bond Market", "US Real Estate Market", "Tech Stocks", "Energy Stocks"
+- Financial Securities: "Apple", "Tesla", "Microsoft", "USD", "Bitcoin", "Gold", "Oil"
+- Economic Indicators: "Consumer Price Index (CPI)", "Unemployment Rate", "GDP Growth Rate", "Federal Deficit"
+
+Input:
+Scores on the Doors: crypto 65.3%, gold 10.1%, stocks 7.7%, HY bonds 4.3%, IG bonds 4.3%, govt bonds 3.4%, 
+oil 3.7%, cash 1.2%, commodities -0.1%, US dollar -2.0% YTD. 
+The Biggest Picture: everyone's new favorite theme...US dollar debasement; US$ -11%
+since Sept, gold >$2k, bitcoin >$30k; right secular theme (deficits, debt, geopolitics),
+US$ in 4th bear market of past 50 years, bullish gold, bearish Euro, bullish international 
+stocks; but pessimism so high right now, and 200bps Fed cuts priced in from June to election, 
+look for US$ trading bounce after Fed ends hiking cycle May 3rd (Chart 2). 
+2024 Tale of the Tape: inflation slowing, Fed hiking cycle over, recession 
+expectations universal, yet UST 30-year can't break below 3.6% (200-dma) because we've already traded 8% to 4% CPI, 
+labor market yet to crack, US govt deficit growing too quickly. The Price is Right: next 6 months it's 
+“recession vs Fed cuts”; best tells who's winning...HY bonds, homebuilders, semiconductors...HYG <73, XHB <70, 
+SOX <2900 recessionary, and if levels hold...it's a no/soft landing.
+
+Output:
+{
+  "entities": [
+    {
+      "name": "crypto",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "gold",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "stocks",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "HY bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "IG bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "govt bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "oil",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "cash",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "commodities",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "US dollar",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "US$",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "gold",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "bitcoin",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "UST 30-year",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "CPI",
+      "type": "Economic Indicators",
+    },
+    {
+      "name": "labor market",
+      "type": "Economic Indicators",
+    },
+    {
+      "name": "US govt deficit",
+      "type": "Economic Indicators",
+    },
+    {
+      "name": "HY bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "homebuilders",
+      "type": "Financial Markets",
+    },
+    {
+      "name": "semiconductors",
+      "type": "Financial Markets",
+    },
+    {
+      "name": "HYG",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "XHB",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "SOX",
+      "type": "Financial Securities",
+    }
+  ]
+}
+
+Input: 
+the latest tech stock resurgence of 2024; NASDAQ +15% since January, Apple >$180, Tesla >$250; 
+driven by strong earnings, innovation, and AI advancements, NASDAQ in its 3rd bull market of the past decade, 
+bullish on semiconductors, cloud computing, and cybersecurity stocks; however, optimism 
+is soaring, and 150bps rate hikes anticipated from July to year-end, expect a potential 
+pullback in tech stocks after the next earnings season concludes in August (Chart 3). 
+Anyway, the bipartisan agreement to suspend the US debt ceiling until 2025 with only limited 
+fiscal restraint (roughly 0.1-0.2% of GDP yoy in 2024 and 2025 compared with the baseline) 
+means that we have avoided another potential significant negative shock to demand.
+On the brightside, the boost to funding that Congress approved late last year for FY23 
+was so large (nearly 10% yoy) that overall discretionary spending is likely to be slightly 
+higher in real terms next year despite the new caps.
+
+Output:
+{
+  "entities": [
+    {
+      "name": "NASDAQ",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "Apple",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "Tesla",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "semiconductors",
+      "type": "Financial Market",
+    },
+    {
+      "name": "cloud computing",
+      "type": "Financial Market",
+    },
+    {
+      "name": "cybersecurity stocks",
+      "type": "Financial Market",
+    },
+    {
+      "name": "tech stocks",
+      "type": "Financial Marker",
+    },
+    {
+      "name": "GDP",
+      "type": "Economic Indicators",
+    },
+    {
+      "name": "FY23",,
+      "type": "Economic Indiciators",
+    }
+  ]
+}
+
+'''
+
+EXTRACT_RELATIONSHIPS_PROMPT = '''
+As a finance expert, your role is to extract relationships between a set of entities. 
+You will receive an excerpt from a financial report, the name of the analyst who authored 
+it, the date the report was published, and a list of entities that have already been extracted. 
+Your task is to identify and output relationships between these entities to construct a 
+knowledge graph for financial document analysis.
+
+Your output should be a JSON object containing a single key, "relationships," 
+which maps to a list of the extracted relationships. These relationships should be 
+formatted as structured JSON objects. Each relationship type must include "from_entity"
+and "to_entity" attributes, using previously extracted entities as guidance, 
+to denote the source and target of the relationship, respectively. Additionally, all 
+relationship types should have "start_date" and "end_date" attributes, formatted as 
+"YYYY-MM-DD". The "source_text" attribute within a relationship should directly quote 
+the description of the relationship from the financial report excerpt.
+
+The relationship types you need to extract include:
+
+[Relationship Type 1: "PREDICTION"]: This relationship captures any forecasts or predictions the analyst makes 
+regarding trends in financial securities, indicators, or market entities. Search for phrases such as "we forecast," 
+"is likely that," or "expect the market to." You must also assess the sentiment of the prediction—whether it is neutral, 
+positive (increase), or negative (decrease)—and include this in the relationship. Map the sentiment to one of the following
+
+- PREDICTED_MOVEMENT_NEUTRAL (neutral sentiment)
+- PREDICTED_MOVEMENT_INCREASE (positive sentiment)
+- PREDICTED_MOVEMENT_DECREASE (negative sentiment)
+
+The "from_entity" should be the analyst's  name(s), while the "to_entity" should be the entity the prediction concerns (utilize the previously extracted entities as guidance). 
+Each prediction must have a defined "start_date" and "end_date" indicating the expected timeframe of the prediction. If 
+the prediction is not explicitly time-bound, omit this relationship. Use the report's publishing date as the "start_date" 
+unless the author specifies a different date. The "end_date" should correspond to when the prediction is anticipated to conclude.
+
+Be sure to include the source excerpt you used to derive each relationship in the "source_text" attribute of the corresponding relationship. 
+Realize that predictions are about the future. Lastly, note that sometimes you will need to infer what entity is being referred to, 
+even if it is not explicitly named in the source sentence. 
+
+Example Input:
+
+Analyst: Helen Brooks
+Publish Date: 2024-06-12
+
+Text Chunk: Labor markets are showing signs of normalization to end 2024; unemployment could drift higher in 2024 while remaining low in historical context.
+Momentum in the job market is starting to wane with slowing payroll growth and modestly rising unemployment, as well as declining quit rates and temporary help. 
+Increased labor force participation and elevated immigration patterns over the past year have added labor supply, while a shortening work week indicates moderating demand for labor. 
+Considering the challenges to add and retain workers coming out of the pandemic, businesses could be more reluctant than normal to shed workers in a slowing economic environment. 
+Even so, less hiring activity could be enough to cause the unemployment rate to tick up to the mid-4% area by the end of next year due to worker churn. 
+Already slowing wage gains should slow further in the context of a softer labor market.
+Anyway, the bipartisan agreement to suspend the US debt ceiling until 2025 with only limited 
+fiscal restraint (roughly 0.1-0.2% of GDP yoy in 2024 and 2025 compared with the baseline) 
+means that we have avoided another potential significant negative shock to demand.
+On the brightside, the boost to funding that Congress approved late last year for FY23 
+was so large (nearly 10% yoy) that overall discretionary spending is likely to be slightly 
+higher in real terms next year despite the new caps.
+
+{
+    "entities": [
+        {
+            "name": "labor markets",
+            "type": "Economic Market",
+        },
+        {
+            "name": "unemployment",
+            "type": "Economic Indicators",
+        },
+        {
+            "name": "job market",
+            "type": "Economic Market",
+        },
+        {
+            "name": "payroll growth",
+            "type": "Economic Indicators",
+        },
+        {
+            "name": "labor force participation",
+            "type": "Economic Indicators",
+        },
+        {
+            "name": "rising unemployment",
+            "type": "Economic Indicators",
+        }
+        {
+            "name": "immigration patterns",
+            "type": "Economic Indicators",
+        },
+        {
+            "name": "hiring activity",
+            "type": "Economic Indicators",
+        },
+        {
+            "name": "worker churn",
+            "type": "Economic Indicators",
+        },
+        {
+          "name": "unemployment rate",
+          "type": "Economic Indicators",
+        },
+        {
+          "name": "discretionary spending",
+          "type": "Economic Indicators",
+        },
+        {
+          "name": "GDP",
+          "type": "Economic Indicators",
+        }
+    ]
+}
+
+
+Example Output:
+{
+    "relationships": [
+        {
+            "PREDICTION": {
+                "from_entity": "Helen Brooks",
+                "to_entity": "labor market",
+                "start_date": "2024-06-12",
+                "end_date": "2024-12-31",
+                "predicted_movement": "PREDICTED_MOVEMENT_NEUTRAL",
+                "source_text": "Labor markets are showing signs of normalization to end 2024;"
+            }
+        },
+        {
+            "PREDICTION": {
+                "from_entity": "Helen Brooks",
+                "to_entity": "unemployment",
+                "startDate": "2024-06-12",
+                "endDate": "2024-12-31",
+                "predicted_movement": "PREDICTED_MOVEMENT_INCREASE",
+                "source_text": "unemployment could drift higher in 2024 while remaining low in historical context."
+            }
+        },
+        {
+            "PREDICTION": {
+                "from_entity": "Helen Brooks",
+                "to_entity": "unemployment rate",
+                "start_date": "2024-06-12",
+                "end_date": "2025-12-31",
+                "predicted_movement": "PREDICTED_MOVEMENT_INCREASE",
+                "source_text": "Even so, less hiring activity could be enough to cause the unemployment rate to tick up to the mid-4% area by the end of next year due to worker churn."
+            }
+        },
+        {
+            "PREDICTION": {
+                "from_entity": "Helen Brooks",
+                "to_entity": "discretionary spending",
+                "start_date": "2024-06-12",
+                "end_date": "2025-12-31",
+                "predicted_movement": "PREDICTED_MOVEMENT_INCREASE",
+                "source_text": "On the brightside, the boost to funding that Congress approved late last year for FY23 was so large (nearly 10% yoy) that overall discretionary spending is likely to be slightly higher in real terms next year despite the new caps."
+            }
+        },
+    ]
+}
+
+Example Input:
+
+Analyst: John Doe, Herb Johnson, Taylor Mason
+
+Publish Date: 2024-01-02
+
+Text Chunk: Scores on the Doors: crypto 65.3%, gold 10.1%, stocks 7.7%, HY bonds 4.3%, IG bonds 4.3%, govt bonds 3.4%,
+oil 3.7%, cash 1.2%, commodities -0.1%, US dollar -2.0% YTD. The Biggest Picture: everyone’s new favorite
+theme...US dollar debasement; US$ -11% since Sept, gold >$2k, bitcoin >$30k; right secular theme (deficits,
+debt, geopolitics), US$ in 4th bear market of past 50 years, bullish gold, oil, Euro, international stocks; but pessimism
+so high right now, and 200bps Fed cuts priced in from June to election, look for US$ trading bounce after Fed
+ends hiking cycle May 3rd (Chart 2). Tale of the Tape: inflation slowing, Fed hiking cycle over, recession
+expectations universal, yet UST 30-year can’t break below 3.6% (200-dma) because we’ve already traded 8% to 4% CPI,
+labor market yet to crack, US govt deficit growing too quickly. The Price is Right: next 6 months it’s
+“recession vs Fed cuts”; best tells who’s winning...HY bonds, homebuilders, semiconductors...HYG <73, XHB <70,
+SOX <2900 recessionary, and if levels hold...it’s a no/soft landing.
+
+{
+  "entities": [
+    {
+      "name": "crypto",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "gold",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "stocks",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "HY bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "IG bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "govt bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "oil",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "cash",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "commodities",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "US dollar",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "US$",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "gold",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "bitcoin",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "UST 30-year",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "CPI",
+      "type": "Economic Indicators",
+    },
+    {
+      "name": "labor market",
+      "type": "Economic Indicators",
+    },
+    {
+      "name": "US govt deficit",
+      "type": "Economic Indicators",
+    },
+    {
+      "name": "HY bonds",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "homebuilders",
+      "type": "Financial Markets",
+    },
+    {
+      "name": "semiconductors",
+      "type": "Financial Markets",
+    },
+    {
+      "name": "HYG",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "XHB",
+      "type": "Financial Securities",
+    },
+    {
+      "name": "SOX",
+      "type": "Financial Securities",
+    }
+  ]
+}
+
+Example Output:
+
+{
+    "relationships": [
+        {
+            "PREDICTION": {
+                "from_entity": "John Doe, Herb Johnson, Taylor Mason",
+                "to_entity": "US$",
+                "start_date": "2024-06-12",
+                "end_date": "2024-11-05",
+                "predicted_movement": "PREDICTED_MOVEMENT_INCREASE",
+                "source_text": "200bps Fed cuts priced in from June to election, look for US$ trading bounce after Fed ends hiking cycle May 3rd."
+            }
+            "PREDICTION": {
+                "from_entity": "John Doe, Herb Johnson, Taylor Mason",
+                "to_entity": "Fed cuts",
+                "start_date": "2024-06-12",
+                "end_date": "2024-11-05",
+                "predicted_movement": "PREDICTED_MOVEMENT_DECREASE",
+                "source_text": "200bps Fed cuts priced in from June to election, look for US$ trading bounce after Fed ends hiking cycle May 3rd."
+            }
+        },
+    ]
+}
+'''

llm_handler/llm_interface.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+from typing import List, Protocol, Type, Dict
+import enum
+
+
+class DefaultEnumMeta(enum.EnumMeta):
+
+    def __call__(cls, value=None, *args, **kwargs) -> DefaultEnumMeta:
+        if value is None:
+            return next(iter(cls))
+        return super().__call__(value, *args, **kwargs)  # type: ignore
+
+
+class LLMInterface(Protocol):
+
+    def get_chat_completion(self, messages: List, model: enum.Enum, temperature: float,
+                            **kwargs) -> str:
+        ...
+
+    def get_text_embedding(self, input: str, model: enum.Enum) -> List[float]:
+        ...

llm_handler/mock_llm_handler.py

@@ -0,0 +1,34 @@
+from typing import List, Optional, Dict
+import enum
+
+import llm_handler.llm_interface as llm_interface
+
+
+class MockLLMHandler(llm_interface.LLMInterface):
+
+    _chat_completion: Optional[List[str]]
+    _text_embedding: Optional[List[float]]
+
+    def __init__(self,
+                 chat_completion: Optional[List[str]] = None,
+                 text_embedding: Optional[List[float]] = None):
+        self._chat_completion = chat_completion
+        self._text_embedding = text_embedding
+
+    def get_chat_completion(self,
+                            messages: List[Dict],
+                            model: Optional[enum.Enum] = None,
+                            temperature: float = 0.2,
+                            **kwargs) -> str:
+        if not self._chat_completion:
+            raise ValueError(f'_chat_completion not set')
+        return self._chat_completion.pop(0)
+
+    def get_text_embedding(
+        self,
+        input: str,
+        model: Optional[enum.Enum] = None,
+    ) -> List[float]:
+        if not self._text_embedding:
+            raise ValueError(f'_text_embedding not set')
+        return self._text_embedding

llm_handler/openai_handler.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+import openai
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
+import os
+from typing import List, Optional, ClassVar
+import enum
+
+from llm_handler.llm_interface import LLMInterface, DefaultEnumMeta
+
+
+class ChatModelVersion(enum.Enum, metaclass=DefaultEnumMeta):
+    GPT_3_5 = 'gpt-3.5-turbo-1106'
+    GPT_4 = 'gpt-4'
+    GPT_4_TURBO = 'gpt-4-1106-preview'
+    GPT_4_O = 'gpt-4o'
+
+
+class EmbeddingModelVersion(enum.Enum, metaclass=DefaultEnumMeta):
+    SMALL_3 = 'text-embedding-3-small'
+    ADA_002 = 'text-embedding-ada-002'
+    LARGE = 'text-embedding-3-large'
+
+
+class OpenAIHandler(LLMInterface):
+
+    _ENV_KEY_NAME: ClassVar[str] = 'OPENAI_API_KEY'
+    _client: openai.Client
+
+    def __init__(self, openai_api_key: Optional[str] = None):
+        _openai_api_key = openai_api_key or os.environ.get(self._ENV_KEY_NAME)
+        if not _openai_api_key:
+            raise ValueError(f'{self._ENV_KEY_NAME} not set')
+        openai.api_key = _openai_api_key
+        self._client = openai.Client()
+
+    def get_chat_completion(  # type: ignore
+            self,
+            messages: List[ChatCompletionMessageParam],
+            model: ChatModelVersion = ChatModelVersion.GPT_4_O,
+            temperature: float = 0.2,
+            **kwargs) -> str:
+        response = self._client.chat.completions.create(model=model.value,
+                                                        messages=messages,
+                                                        temperature=temperature,
+                                                        **kwargs)
+        responses: List[str] = []
+        for choice in response.choices:
+            if choice.finish_reason != 'stop' or not choice.message.content:
+                raise ValueError(f'Choice did not complete correctly: {choice}')
+            responses.append(choice.message.content)
+        if len(responses) != 1:
+            raise ValueError(f'Expected one response, got {len(responses)}: {responses}')
+        return responses[0]
+
+    def get_text_embedding(  # type: ignore
+            self, input: str, model: EmbeddingModelVersion) -> List[float]:
+        response = self._client.embeddings.create(model=model.value,
+                                                  encoding_format='float',
+                                                  input=input)
+        if not response.data:
+            raise ValueError(f'No embedding in response: {response}')
+        elif len(response.data) != 1:
+            raise ValueError(f'More than one embedding in response: {response}')
+        return response.data[0].embedding

llm_handler/openai_handler_test.py

@@ -0,0 +1,69 @@
+import os
+from typing import List
+from unittest import mock
+import pytest
+import openai
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
+
+import llm_handler.openai_handler as openai_handler
+
+
+# Clear openai.api_key before each test
+@pytest.fixture(autouse=True)
+def test_env_setup():
+    openai.api_key = None
+
+
+# This allows us to clear the OPENAI_API_KEY before any test we want
+@pytest.fixture()
+def mock_settings_env_vars():
+    with mock.patch.dict(os.environ, clear=True):
+        yield
+
+
+# Define some constants for our tests
+_TEST_KEY: str = "TEST_KEY"
+_TEST_MESSAGE: str = "Hello how are you?"
+_TEST_MESSAGE_FOR_CHAT_COMPLETION: List[ChatCompletionMessageParam] = [
+    {
+        "role": "system", "content": "You are serving as a en endpoint to verify a test."
+    }, {
+        "role": "user", "content": "Respond with something to help us verify our code is working."
+    }
+]
+_TEXT_EMBEDDING_SMALL_3_LENGTH: int = 1536
+
+
+def test_init_without_key(mock_settings_env_vars):
+    # Ensure key not set as env var and openai.api_key not set
+    with pytest.raises(KeyError):
+        os.environ[openai_handler.OpenAIHandler._ENV_KEY_NAME]
+    assert openai.api_key == None
+    # Ensure proper exception raised when instantiating handler without key as param or env var
+    with pytest.raises(ValueError) as excinfo:
+        openai_handler.OpenAIHandler()
+    assert f'{openai_handler.OpenAIHandler._ENV_KEY_NAME} not set' in str(excinfo.value)
+    assert openai.api_key == None
+
+
+def test_init_with_key_as_param():
+    # Ensure key is set as env var, key value is unique from _TEST_KEY, and openai.api_key not set
+    assert not os.environ[openai_handler.OpenAIHandler._ENV_KEY_NAME] == _TEST_KEY
+    assert openai.api_key == None
+    # Ensure successful instantiation and openai.api_key properly set
+    handler = openai_handler.OpenAIHandler(openai_api_key=_TEST_KEY)
+    assert isinstance(handler, openai_handler.OpenAIHandler)
+    assert openai.api_key == _TEST_KEY
+
+
+def test_init_with_key_as_env_var(mock_settings_env_vars):
+    # Ensure key not set as env var and openai.api_key not set
+    with pytest.raises(KeyError):
+        os.environ[openai_handler.OpenAIHandler._ENV_KEY_NAME]
+    assert openai.api_key == None
+    # Set key as env var
+    os.environ.setdefault(openai_handler.OpenAIHandler._ENV_KEY_NAME, _TEST_KEY)
+    # Ensure successful instantiation and openai.api_key properly set
+    handler = openai_handler.OpenAIHandler()
+    assert isinstance(handler, openai_handler.OpenAIHandler)
+    assert openai.api_key == _TEST_KEY

proto/chunk_pb2.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: proto/chunk.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x11proto/chunk.proto\x12\x06\x66inllm\"D\n\x08\x44ocument\x12\x11\n\tfile_path\x18\x01 \x01(\t\x12\x0f\n\x07\x61uthors\x18\x02 \x01(\t\x12\x14\n\x0cpublish_date\x18\x03 \x01(\t\"\xbe\x01\n\x05\x43hunk\x12\x10\n\x08\x63hunk_id\x18\x01 \x01(\t\x12\x12\n\nchunk_text\x18\x02 \x01(\t\x12%\n\nchunk_type\x18\x03 \x01(\x0e\x32\x11.finllm.ChunkType\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\x03\x12\x19\n\x0fparent_chunk_id\x18\x05 \x01(\tH\x00\x12$\n\x08\x64ocument\x18\x06 \x01(\x0b\x32\x10.finllm.DocumentH\x00\x42\x12\n\x10parent_reference*k\n\tChunkType\x12\x16\n\x12\x43HUNK_TYPE_UNKNOWN\x10\x00\x12\x17\n\x13\x43HUNK_TYPE_SENTENCE\x10\x01\x12\x18\n\x14\x43HUNK_TYPE_PARAGRAPH\x10\x02\x12\x13\n\x0f\x43HUNK_TYPE_PAGE\x10\x03\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'proto.chunk_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_CHUNKTYPE']._serialized_start=292
+  _globals['_CHUNKTYPE']._serialized_end=399
+  _globals['_DOCUMENT']._serialized_start=29
+  _globals['_DOCUMENT']._serialized_end=97
+  _globals['_CHUNK']._serialized_start=100
+  _globals['_CHUNK']._serialized_end=290
+# @@protoc_insertion_point(module_scope)

proto/chunk_pb2.pyi

@@ -0,0 +1,43 @@
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class ChunkType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    CHUNK_TYPE_UNKNOWN: _ClassVar[ChunkType]
+    CHUNK_TYPE_SENTENCE: _ClassVar[ChunkType]
+    CHUNK_TYPE_PARAGRAPH: _ClassVar[ChunkType]
+    CHUNK_TYPE_PAGE: _ClassVar[ChunkType]
+CHUNK_TYPE_UNKNOWN: ChunkType
+CHUNK_TYPE_SENTENCE: ChunkType
+CHUNK_TYPE_PARAGRAPH: ChunkType
+CHUNK_TYPE_PAGE: ChunkType
+
+class Document(_message.Message):
+    __slots__ = ("file_path", "authors", "publish_date")
+    FILE_PATH_FIELD_NUMBER: _ClassVar[int]
+    AUTHORS_FIELD_NUMBER: _ClassVar[int]
+    PUBLISH_DATE_FIELD_NUMBER: _ClassVar[int]
+    file_path: str
+    authors: str
+    publish_date: str
+    def __init__(self, file_path: _Optional[str] = ..., authors: _Optional[str] = ..., publish_date: _Optional[str] = ...) -> None: ...
+
+class Chunk(_message.Message):
+    __slots__ = ("chunk_id", "chunk_text", "chunk_type", "chunk_index", "parent_chunk_id", "document")
+    CHUNK_ID_FIELD_NUMBER: _ClassVar[int]
+    CHUNK_TEXT_FIELD_NUMBER: _ClassVar[int]
+    CHUNK_TYPE_FIELD_NUMBER: _ClassVar[int]
+    CHUNK_INDEX_FIELD_NUMBER: _ClassVar[int]
+    PARENT_CHUNK_ID_FIELD_NUMBER: _ClassVar[int]
+    DOCUMENT_FIELD_NUMBER: _ClassVar[int]
+    chunk_id: str
+    chunk_text: str
+    chunk_type: ChunkType
+    chunk_index: int
+    parent_chunk_id: str
+    document: Document
+    def __init__(self, chunk_id: _Optional[str] = ..., chunk_text: _Optional[str] = ..., chunk_type: _Optional[_Union[ChunkType, str]] = ..., chunk_index: _Optional[int] = ..., parent_chunk_id: _Optional[str] = ..., document: _Optional[_Union[Document, _Mapping]] = ...) -> None: ...

proto/entity_pb2.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: proto/entity.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12proto/entity.proto\x12\x06\x66inllm\"0\n\x06\x45ntity\x12\x11\n\tentity_id\x18\x01 \x01(\t\x12\x13\n\x0b\x65ntity_name\x18\x02 \x01(\t\"\x99\x01\n\x0cRelationship\x12\x17\n\x0frelationship_id\x18\x01 \x01(\t\x12\x12\n\nstart_date\x18\x02 \x01(\t\x12\x10\n\x08\x65nd_date\x18\x03 \x01(\t\x12\x13\n\x0bsource_text\x18\x04 \x01(\t\x12\x35\n\x12predicted_movement\x18\x05 \x01(\x0e\x32\x19.finllm.PredictedMovement\"\x88\x01\n\x12\x45ntityRelationship\x12#\n\x0b\x66rom_entity\x18\x01 \x01(\x0b\x32\x0e.finllm.Entity\x12*\n\x0crelationship\x18\x02 \x01(\x0b\x32\x14.finllm.Relationship\x12!\n\tto_entity\x18\x03 \x01(\x0b\x32\x0e.finllm.Entity\"P\n\x14\x45ntityKnowledgeGraph\x12\x38\n\x14\x65ntity_relationships\x18\x01 \x03(\x0b\x32\x1a.finllm.EntityRelationship*\x99\x01\n\x11PredictedMovement\x12\"\n\x1ePREDICTED_MOVEMENT_UNSPECIFIED\x10\x00\x12\x1e\n\x1aPREDICTED_MOVEMENT_NEUTRAL\x10\x01\x12\x1f\n\x1bPREDICTED_MOVEMENT_INCREASE\x10\x02\x12\x1f\n\x1bPREDICTED_MOVEMENT_DECREASE\x10\x03\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'proto.entity_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_PREDICTEDMOVEMENT']._serialized_start=458
+  _globals['_PREDICTEDMOVEMENT']._serialized_end=611
+  _globals['_ENTITY']._serialized_start=30
+  _globals['_ENTITY']._serialized_end=78
+  _globals['_RELATIONSHIP']._serialized_start=81
+  _globals['_RELATIONSHIP']._serialized_end=234
+  _globals['_ENTITYRELATIONSHIP']._serialized_start=237
+  _globals['_ENTITYRELATIONSHIP']._serialized_end=373
+  _globals['_ENTITYKNOWLEDGEGRAPH']._serialized_start=375
+  _globals['_ENTITYKNOWLEDGEGRAPH']._serialized_end=455
+# @@protoc_insertion_point(module_scope)

proto/entity_pb2.pyi

@@ -0,0 +1,56 @@
+from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class PredictedMovement(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    PREDICTED_MOVEMENT_UNSPECIFIED: _ClassVar[PredictedMovement]
+    PREDICTED_MOVEMENT_NEUTRAL: _ClassVar[PredictedMovement]
+    PREDICTED_MOVEMENT_INCREASE: _ClassVar[PredictedMovement]
+    PREDICTED_MOVEMENT_DECREASE: _ClassVar[PredictedMovement]
+PREDICTED_MOVEMENT_UNSPECIFIED: PredictedMovement
+PREDICTED_MOVEMENT_NEUTRAL: PredictedMovement
+PREDICTED_MOVEMENT_INCREASE: PredictedMovement
+PREDICTED_MOVEMENT_DECREASE: PredictedMovement
+
+class Entity(_message.Message):
+    __slots__ = ("entity_id", "entity_name")
+    ENTITY_ID_FIELD_NUMBER: _ClassVar[int]
+    ENTITY_NAME_FIELD_NUMBER: _ClassVar[int]
+    entity_id: str
+    entity_name: str
+    def __init__(self, entity_id: _Optional[str] = ..., entity_name: _Optional[str] = ...) -> None: ...
+
+class Relationship(_message.Message):
+    __slots__ = ("relationship_id", "start_date", "end_date", "source_text", "predicted_movement")
+    RELATIONSHIP_ID_FIELD_NUMBER: _ClassVar[int]
+    START_DATE_FIELD_NUMBER: _ClassVar[int]
+    END_DATE_FIELD_NUMBER: _ClassVar[int]
+    SOURCE_TEXT_FIELD_NUMBER: _ClassVar[int]
+    PREDICTED_MOVEMENT_FIELD_NUMBER: _ClassVar[int]
+    relationship_id: str
+    start_date: str
+    end_date: str
+    source_text: str
+    predicted_movement: PredictedMovement
+    def __init__(self, relationship_id: _Optional[str] = ..., start_date: _Optional[str] = ..., end_date: _Optional[str] = ..., source_text: _Optional[str] = ..., predicted_movement: _Optional[_Union[PredictedMovement, str]] = ...) -> None: ...
+
+class EntityRelationship(_message.Message):
+    __slots__ = ("from_entity", "relationship", "to_entity")
+    FROM_ENTITY_FIELD_NUMBER: _ClassVar[int]
+    RELATIONSHIP_FIELD_NUMBER: _ClassVar[int]
+    TO_ENTITY_FIELD_NUMBER: _ClassVar[int]
+    from_entity: Entity
+    relationship: Relationship
+    to_entity: Entity
+    def __init__(self, from_entity: _Optional[_Union[Entity, _Mapping]] = ..., relationship: _Optional[_Union[Relationship, _Mapping]] = ..., to_entity: _Optional[_Union[Entity, _Mapping]] = ...) -> None: ...
+
+class EntityKnowledgeGraph(_message.Message):
+    __slots__ = ("entity_relationships",)
+    ENTITY_RELATIONSHIPS_FIELD_NUMBER: _ClassVar[int]
+    entity_relationships: _containers.RepeatedCompositeFieldContainer[EntityRelationship]
+    def __init__(self, entity_relationships: _Optional[_Iterable[_Union[EntityRelationship, _Mapping]]] = ...) -> None: ...

proto/pdf_document_pb2.py

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: proto/pdf_document.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x18proto/pdf_document.proto\x12\x06\x66inllm\"!\n\x0cPdfDocumentD\x12\x11\n\tfile_path\x18\x01 \x01(\tb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'proto.pdf_document_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_PDFDOCUMENTD']._serialized_start=36
+  _globals['_PDFDOCUMENTD']._serialized_end=69
+# @@protoc_insertion_point(module_scope)

proto/pdf_document_pb2.pyi

@@ -0,0 +1,11 @@
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Optional as _Optional
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class PdfDocumentD(_message.Message):
+    __slots__ = ("file_path",)
+    FILE_PATH_FIELD_NUMBER: _ClassVar[int]
+    file_path: str
+    def __init__(self, file_path: _Optional[str] = ...) -> None: ...

query_pipeline/evaluation_engine.py

@@ -0,0 +1,219 @@
+import json
+from typing import Optional
+
+from proto.entity_pb2 import PredictedMovement
+from tabulate import tabulate
+
+from domain.entity_d import (
+    EntityD,
+    EntityKnowledgeGraphD,
+    EntityRelationshipD,
+    RelationshipD,
+)
+from llm_handler.openai_handler import (
+    ChatCompletionMessageParam,
+    ChatModelVersion,
+    OpenAIHandler,
+)
+from utils.dates import parse_date
+
+FUZZY_MATCH_ENTITIES_PROMPT = '''
+You are an expert in financial analysis. You will be given a two lists of entities. Your task is to output a semantic mapping from the entities in list A to the entities in list B. This means an entity in list A that is semantically similar to an entity in list B should be mapped together. If there is no reasonable semantic match for an entity in list A, output an empty string. Output should be in the format of a JSON object. Ensure the entity keys are in the same order as the input list A.
+
+Input:
+List A: ["BofA", "Bank of Amerca Corp" "GDP", "Inflation", "Yen"]
+List B: ["Bank of America", "inflation", "Gross Domestic Product", "oil"]
+
+Output:
+{
+    "BofA": "Bank of America",
+    "Bank of America Corp": "Bank of America",
+    "GDP": "Gross Domestic Product",
+    "Inflation": "inflation",
+    "Yen": ""
+}
+'''
+
+
+class EvaluationEngine:
+
+    _handler: OpenAIHandler
+    _MODEL_VERSION: ChatModelVersion = ChatModelVersion.GPT_4_O
+    _TEMPERATURE: float = 0.2
+
+    def __init__(self,
+                 ground_truth_kg: EntityKnowledgeGraphD,
+                 openai_handler: Optional[OpenAIHandler] = None,
+                 model_version: Optional[ChatModelVersion] = None):
+        self._handler = openai_handler or OpenAIHandler()
+        self._model_version = model_version or self._MODEL_VERSION
+
+        # setup adjacency list representation of ground truth knowledge graph
+        self.kg: dict[str, list[EntityRelationshipD]] = {}
+        for entity_relationship in ground_truth_kg.entity_relationships:
+            to_entity_name = entity_relationship.to_entity.entity_name
+            relationships = self.kg.get(to_entity_name, [])
+            relationships.append(entity_relationship)
+            self.kg[to_entity_name] = relationships
+
+    def _get_thesis_to_gt_entity_map(self, thesis_kg: EntityKnowledgeGraphD) -> dict[str, str]:
+        thesis_entities = []
+        for entity_relationship in thesis_kg.entity_relationships:
+            thesis_entities.append(entity_relationship.to_entity.entity_name)
+
+        # LLM call to return out the matched entities
+        messages: list[ChatCompletionMessageParam] = [
+            {
+                "role": "system", "content": FUZZY_MATCH_ENTITIES_PROMPT
+            }, {
+                "role": "user",
+                "content": f"List A: {thesis_entities}\nList B: {list(self.kg.keys())}"
+            }
+        ]
+        completion_text = self._handler.get_chat_completion(messages=messages,
+                                                            model=self._model_version,
+                                                            temperature=self._TEMPERATURE,
+                                                            response_format={"type": "json_object"})
+
+        thesis_to_gt_entity_mapping: dict[str, str] = json.loads(completion_text)
+
+        return thesis_to_gt_entity_mapping
+
+    def _get_relationships_matching_timeperiod(
+            self, gt_kg_to_node: str, relationship: RelationshipD) -> list[EntityRelationshipD]:
+        matching_relationships = []
+        thesis_relationship_start = parse_date(relationship.start_date)
+        thesis_relationship_end = parse_date(relationship.end_date)
+        for gt_relationship in self.kg[gt_kg_to_node]:
+            gt_relationship_start = parse_date(gt_relationship.relationship.start_date)
+            gt_relationship_end = parse_date(gt_relationship.relationship.end_date)
+
+            if (gt_relationship_start <= thesis_relationship_start <= gt_relationship_end and \
+                gt_relationship_start <= thesis_relationship_end <= gt_relationship_end):
+                # thesis relationship timeframe and gt relationship timeframe overlap
+                matching_relationships.append(gt_relationship)
+
+        return matching_relationships
+
+    def evaluate_thesis(
+        self, thesis_kg: EntityKnowledgeGraphD
+    ) -> list[tuple[EntityRelationshipD, bool, Optional[EntityRelationshipD]]]:
+        thesis_to_kg_map = self._get_thesis_to_gt_entity_map(thesis_kg)
+        results = []
+        for thesis_relationship in thesis_kg.entity_relationships:
+            thesis_to_node = thesis_relationship.to_entity.entity_name
+            kg_node = thesis_to_kg_map[thesis_to_node]
+            if not kg_node:  # no matching entity in KG
+                results.append((thesis_relationship, False, None))
+                continue
+
+            matching_relationships = self._get_relationships_matching_timeperiod(
+                kg_node, thesis_relationship.relationship)
+
+            for entity_relationship in matching_relationships:
+                if entity_relationship.relationship.predicted_movement == thesis_relationship.relationship.predicted_movement:
+                    results.append((thesis_relationship, True, entity_relationship))
+                else:
+                    results.append((thesis_relationship, False, entity_relationship))
+            if len(matching_relationships) == 0:
+                results.append((thesis_relationship, False, None))
+
+        return results
+
+    def evaluate_and_display_thesis(self, thesis_kg: EntityKnowledgeGraphD):
+        results = self.evaluate_thesis(thesis_kg)
+
+        int_to_str = {1: "Neutral", 2: 'Increase', 3: 'Decrease'}
+
+        headers = ["Thesis Claim", "Supported by KG", "Related KG Relationship"]
+        table_data = []
+        for triplet in results:
+            claim_entity = triplet[0].to_entity.entity_name
+            claim_movement = int_to_str[triplet[0].relationship.predicted_movement]
+            claim = f'{claim_entity} {claim_movement}'
+            if triplet[2]:
+                evidence = int_to_str[triplet[2].relationship.predicted_movement]
+                evidence += f' ({triplet[2].from_entity.entity_name}) '
+            else:
+                evidence = "No evidence in KG"
+            table_data.append([claim, triplet[1], evidence])
+        return tabulate(table_data, tablefmt="html", headers=headers)
+
+
+if __name__ == '__main__':
+    # TODO: extract the cases into pytest tests
+    kg = EntityKnowledgeGraphD(entity_relationships=[
+        EntityRelationshipD(from_entity=EntityD(entity_id='3', entity_name="analyst A"),
+                            relationship=RelationshipD(
+                                relationship_id='2',
+                                start_date='2021-01-01',
+                                end_date='2024-12-31',
+                                source_text='',
+                                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_INCREASE),
+                            to_entity=EntityD(entity_id='1', entity_name="GDP")),
+        EntityRelationshipD(from_entity=EntityD(entity_id='5', entity_name="analyst B"),
+                            relationship=RelationshipD(
+                                relationship_id='3',
+                                start_date='2021-01-01',
+                                end_date='2021-12-31',
+                                source_text='',
+                                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_DECREASE),
+                            to_entity=EntityD(entity_id='1', entity_name="GDP")),
+        EntityRelationshipD(from_entity=EntityD(entity_id='7', entity_name="analyst C"),
+                            relationship=RelationshipD(
+                                relationship_id='4',
+                                start_date='2021-01-01',
+                                end_date='2021-12-31',
+                                source_text='',
+                                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL),
+                            to_entity=EntityD(entity_id='1', entity_name="GDP")),
+        EntityRelationshipD(from_entity=EntityD(entity_id='9', entity_name="analyst D"),
+                            relationship=RelationshipD(
+                                relationship_id='5',
+                                start_date='2021-01-01',
+                                end_date='2021-12-31',
+                                source_text='',
+                                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL),
+                            to_entity=EntityD(entity_id='10', entity_name="USD")),
+        EntityRelationshipD(  # out of time range for thesis
+            from_entity=EntityD(entity_id='9', entity_name="analyst E"),
+            relationship=RelationshipD(
+                relationship_id='5',
+                start_date='2024-01-01',
+                end_date='2024-12-31',
+                source_text='',
+                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_NEUTRAL),
+            to_entity=EntityD(entity_id='10', entity_name="USD")),
+    ])
+
+    thesis_claims = [
+        EntityRelationshipD(from_entity=EntityD(entity_id='2', entity_name="user"),
+                            relationship=RelationshipD(
+                                relationship_id='1',
+                                start_date='2021-01-01',
+                                end_date='2021-12-31',
+                                source_text='',
+                                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_INCREASE),
+                            to_entity=EntityD(entity_id='1', entity_name="Gross Domestic Product")),
+        EntityRelationshipD(from_entity=EntityD(entity_id='2', entity_name="user"),
+                            relationship=RelationshipD(
+                                relationship_id='1',
+                                start_date='2021-01-01',
+                                end_date='2021-12-31',
+                                source_text='',
+                                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_INCREASE),
+                            to_entity=EntityD(entity_id='1', entity_name="US$")),
+        EntityRelationshipD(from_entity=EntityD(entity_id='2', entity_name="user"),
+                            relationship=RelationshipD(
+                                relationship_id='1',
+                                start_date='2021-01-01',
+                                end_date='2021-12-31',
+                                source_text='',
+                                predicted_movement=PredictedMovement.PREDICTED_MOVEMENT_INCREASE),
+                            to_entity=EntityD(entity_id='1', entity_name="Yen")),
+    ]
+    thesis = EntityKnowledgeGraphD(entity_relationships=thesis_claims)
+
+    eval_engine = EvaluationEngine(kg)
+
+    eval_engine.evaluate_and_display_thesis(thesis)

query_pipeline/thesis_extractor.py

@@ -0,0 +1,18 @@
+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)

requirements.txt

@@ -0,0 +1,39 @@
+annotated-types==0.7.0
+anyio==4.4.0
+certifi==2024.6.2
+distro==1.9.0
+exceptiongroup==1.2.1
+grpcio==1.64.1
+grpcio-tools==1.64.1
+h11==0.14.0
+httpcore==1.0.5
+httpx==0.27.0
+idna==3.7
+importlib_metadata==7.1.0
+iniconfig==2.0.0
+neo4j==5.21.0
+nodeenv==1.9.1
+numpy==1.26.4
+openai==1.33.0
+packaging==24.0
+pandas==2.2.2
+platformdirs==4.2.2
+pluggy==1.5.0
+protobuf==5.27.1
+pydantic==2.7.3
+pydantic_core==2.18.4
+PyMuPDF==1.24.5
+PyMuPDFb==1.24.3
+pyright==1.1.366
+pytest==8.2.2
+python-dateutil==2.9.0.post0
+pytz==2024.1
+six==1.16.0
+sniffio==1.3.1
+tabulate==0.9.0
+tomli==2.0.1
+tqdm==4.66.4
+typing_extensions==4.12.2
+tzdata==2024.1
+yapf==0.40.2
+zipp==3.19.2

storage/domain_dao.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+from typing import List, Type, Protocol, TypeVar, Dict, Set
+import os
+import json
+import uuid
+
+from domain.domain_protocol import DomainProtocol
+
+DomainT = TypeVar('DomainT', bound=DomainProtocol)
+MAP_BIN = os.path.join(os.path.dirname(os.path.abspath(os.path.dirname(__file__))), ".bin", "maps")
+
+
+class DomainDAO(Protocol[DomainT]):
+
+    def insert(self, domain_objs: List[DomainT]):
+        ...
+
+    def read_by_id(self, domain_id: str) -> DomainT:
+        ...
+
+    def read_all(self) -> Set[DomainT]:
+        ...
+
+
+class InMemDomainDAO(DomainDAO[DomainT]):
+
+    _id_to_domain_obj: Dict[str, DomainT]
+
+    def __init__(self):
+        self._id_to_domain_obj = {}
+
+    def insert(self, domain_objs: List[DomainT]):
+        new_id_to_domain_obj = {domain_obj.id: domain_obj for domain_obj in domain_objs}
+        if len(new_id_to_domain_obj) != len(domain_objs):
+            raise ValueError("Duplicate IDs exist within incoming domain_objs")
+        if duplicate_ids := set(new_id_to_domain_obj.keys()) & set(self._id_to_domain_obj.keys()):
+            raise ValueError(f"Duplicate ids exist in DB: {duplicate_ids}")
+        self._id_to_domain_obj.update(new_id_to_domain_obj)
+
+    def read_by_id(self, domain_id: str) -> DomainT:
+        if domain_obj := self._id_to_domain_obj.get(domain_id):
+            return domain_obj
+        raise ValueError(f"Domain obj with id {domain_id} not found")
+
+    def read_all(self) -> Set[DomainT]:
+        return set(self._id_to_domain_obj.values())
+
+    @classmethod
+    def load_from_file(cls, file_path: str, domain_cls: Type[DomainT]) -> InMemDomainDAO[DomainT]:
+        if not os.path.isfile(file_path):
+            raise ValueError(f"File not found: {file_path}")
+        with open(file_path, 'r') as f:
+            domain_objs = [domain_cls.from_json(line) for line in f]
+        dao = cls()
+        dao.insert(domain_objs)
+        return dao
+
+    def save_to_file(self, file_path: str):
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        domain_jsons = [domain_obj.to_json() for domain_obj in self._id_to_domain_obj.values()]
+        with open(file_path, 'w') as f:
+            f.write('\n'.join(domain_jsons) + '\n')
+
+
+class CacheDomainDAO(DomainDAO[DomainT]):
+
+    _id_to_domain_obj: Dict[str, DomainT]
+    _save_path: str
+
+    def __init__(self, save_path: str, domain_cls: Type[DomainT]):
+        self._id_to_domain_obj = {}
+        self._save_path = os.path.join(MAP_BIN, save_path)
+        self._load_cache(domain_cls)
+
+    def __enter__(self):
+        return self
+
+    def __call__(self, element: DomainT) -> DomainT:
+        self.insert([element])
+        return element
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._save_cache()
+
+    def set(self, element: DomainT) -> uuid.UUID:
+        id = uuid.uuid4()
+        self._id_to_domain_obj[str(id)] = element
+        self._save_cache()
+        return id
+
+    def _save_cache(self):
+        os.makedirs(MAP_BIN, exist_ok=True)
+        cache = {}
+        if os.path.isfile(self._save_path):
+            with open(self._save_path, 'r') as f:
+                cache = json.load(f)
+        domain_json_map = {
+            id: domain_obj.to_json()
+            for id, domain_obj in self._id_to_domain_obj.items()
+        }
+        cache.update(domain_json_map)
+        with open(self._save_path, 'w') as f:
+            json.dump(cache, f, indent=4)
+
+    def _load_cache(self, domain_cls: Type[DomainT]):
+        if not os.path.isfile(self._save_path):
+            return
+        with open(self._save_path, 'r') as f:
+            domain_json_map = json.load(f)
+        for id, domain_json in domain_json_map.items():
+            self._id_to_domain_obj[id] = domain_cls.from_json(domain_json)
+
+    def read_by_id(self, domain_id: str) -> DomainT:
+        if domain_obj := self._id_to_domain_obj.get(domain_id):
+            return domain_obj
+        raise ValueError(f"Domain obj with id {domain_id} not found")
+
+    def read_all(self) -> Set[DomainT]:
+        return set(self._id_to_domain_obj.values())
+
+    def insert(self, domain_objs: List[DomainT]):
+        new_id_to_domain_obj = {domain_obj.id: domain_obj for domain_obj in domain_objs}
+        if len(new_id_to_domain_obj) != len(domain_objs):
+            raise ValueError("Duplicate IDs exist within incoming domain_objs")
+        if duplicate_ids := set(new_id_to_domain_obj.keys()) & set(self._id_to_domain_obj.keys()):
+            raise ValueError(f"Duplicate ids exist in DB: {duplicate_ids}")
+        self._id_to_domain_obj.update(new_id_to_domain_obj)
+        self._save_cache()

storage/domain_dao_test.py

@@ -0,0 +1,71 @@
+import logging
+import unittest
+import uuid
+import os
+from storage.domain_dao import InMemDomainDAO
+from domain.domain_protocol_test import TimestampTestD
+
+
+class InMemDomainDAOTest(unittest.TestCase):
+
+    def setUp(self):
+        self.in_mem_dao: InMemDomainDAO = InMemDomainDAO[TimestampTestD]()
+        self.timestamp_d: TimestampTestD = TimestampTestD(nanos=1)
+
+    def test_insert_domain_obj(self):
+        self.in_mem_dao.insert([self.timestamp_d])
+        expected_map = {self.timestamp_d.id: self.timestamp_d}
+        self.assertEqual(self.in_mem_dao._id_to_domain_obj, expected_map)
+
+    def test_insert_domain_obj_raise_on_duplicate_id_in_db(self):
+        self.in_mem_dao.insert([self.timestamp_d])
+        with self.assertRaises(ValueError) as context:
+            self.in_mem_dao.insert([self.timestamp_d])
+        self.assertIn("Duplicate ids exist in DB", str(context.exception))
+
+    def test_insert_domain_obj_raise_on_duplicate_id_arguements(self):
+        with self.assertRaises(ValueError) as context:
+            self.in_mem_dao.insert([self.timestamp_d, self.timestamp_d])
+        self.assertIn("Duplicate IDs exist within incoming domain_objs", str(context.exception))
+
+    def test_read_by_id(self):
+        self.in_mem_dao.insert([self.timestamp_d])
+        timestamp_d = self.in_mem_dao.read_by_id(self.timestamp_d.id)
+        self.assertEqual(timestamp_d, self.timestamp_d)
+
+    def test_read_by_id_raise_not_found(self):
+        with self.assertRaises(ValueError):
+            self.in_mem_dao.read_by_id(self.timestamp_d.id)
+
+    def test_read_all(self):
+        timestamp_d_b = TimestampTestD(2)
+        self.in_mem_dao.insert([self.timestamp_d, timestamp_d_b])
+        expected_timestamps = {self.timestamp_d, timestamp_d_b}
+        self.assertEqual(expected_timestamps, self.in_mem_dao.read_all())
+
+    def test_load_from_file(self):
+        file_path = f".bin/{uuid.uuid4()}.jsonl"
+        with open(file_path, 'w') as f:
+            f.write(self.timestamp_d.to_json() + '\n')
+        dao = InMemDomainDAO[TimestampTestD].load_from_file(file_path, TimestampTestD)
+        os.remove(file_path)
+        self.assertEqual({self.timestamp_d}, dao.read_all())
+
+    def test_load_from_file_fail_not_found(self):
+        with self.assertRaises(ValueError):
+            _ = InMemDomainDAO[TimestampTestD].load_from_file("file_path", TimestampTestD)
+
+    def test_save_to_file(self):
+        file_path = f".bin/{uuid.uuid4()}.jsonl"
+        self.in_mem_dao.insert([self.timestamp_d])
+        self.in_mem_dao.save_to_file(file_path)
+        created_dao = InMemDomainDAO[TimestampTestD].load_from_file(file_path, TimestampTestD)
+        os.remove(file_path)
+        self.assertEqual(self.in_mem_dao.read_all(), created_dao.read_all())
+
+
+#TODO: Add test for CacheDomainDAO
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.DEBUG)
+    unittest.main()

storage/neo4j_dao.py

@@ -0,0 +1,102 @@
+import logging
+import os
+from typing import Optional
+
+import neo4j
+
+from domain.entity_d import (
+    EntityD,
+    EntityKnowledgeGraphD,
+    EntityRelationshipD,
+    RelationshipD,
+)
+
+
+class Neo4jError(Exception):
+    ...
+
+
+class Neo4jDomainDAO:
+    """ To be used with a context manager to ensure the connection is closed after use. """
+
+    def __enter__(self):
+        uri = os.environ.get("NEO4J_URI", "")
+        user = os.environ.get("NEO4J_USER", "")
+        password = os.environ.get("NEO4J_PASSWORD", "")
+
+        if not uri:
+            raise ValueError("NEO4J_URI environment variable not set")
+        if not user:
+            raise ValueError("NEO4J_USER environment variable not set")
+        if not password:
+            raise ValueError("NEO4J_PASSWORD environment variable not set")
+
+        try:
+            self.driver = neo4j.GraphDatabase.driver(uri, auth=(user, password))
+            self.driver.verify_connectivity()
+        except Exception as e:
+            logging.error(f"Failed to connect to Neo4j: {e}")
+            raise Neo4jError("Failed to connect to Neo4j")
+
+        return self
+
+    def insert(self, knowledge_graph: EntityKnowledgeGraphD, pdf_file: str = ""):
+        for entity_relationship in knowledge_graph.entity_relationships:
+            create_cmds = entity_relationship.neo4j_create_cmds
+            create_cmds_args = entity_relationship.neo4j_create_args
+
+            for create_cmd, args in zip(create_cmds, create_cmds_args):
+                args['pdf_file'] = pdf_file
+                try:
+                    self.driver.execute_query(
+                        create_cmd,  # type: ignore
+                        parameters_=args,  # type: ignore
+                        database_='neo4j')  # type: ignore
+                except Exception as e:
+                    logging.warning(
+                        f"Failed to insert entity relationship: {entity_relationship} due to {e}")
+
+    def query(self, query, query_args):
+        return self.driver.execute_query(query, parameters_=query_args,
+                                         database_='neo4j')  # type: ignore
+
+    def get_knowledge_graph(self) -> Optional[EntityKnowledgeGraphD]:
+        records = []  #list[dict[str, Neo4jDict]]
+        try:
+            records, _, _ = self.driver.execute_query("MATCH (from:Entity) -[r:Relationship]-> (to:Entity) RETURN from, properties(r), to", database_='neo4j')  # type: ignore
+        except Exception as e:
+            logging.exception(e)
+            return None
+
+        entity_relationships = []
+        for record in records:
+            er_dict = record.data()
+
+            from_args = er_dict['from']
+            from_entity = EntityD(entity_id='', entity_name=from_args['name'])
+
+            to_args = er_dict['to']
+            to_entity = EntityD(entity_id='', entity_name=to_args['name'])
+
+            relationship_args = er_dict['properties(r)']
+            relationship = RelationshipD(relationship_id='',
+                                         start_date=relationship_args['start_date'],
+                                         end_date=relationship_args['end_date'],
+                                         source_text=relationship_args['source_text'],
+                                         predicted_movement=RelationshipD.from_string(
+                                             relationship_args['predicted_movement']))
+
+            entity_relationships.append(
+                EntityRelationshipD(from_entity=from_entity,
+                                    relationship=relationship,
+                                    to_entity=to_entity))
+
+        return EntityKnowledgeGraphD(entity_relationships=entity_relationships)
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        if traceback:
+            logging.error("Neo4jDomainDAO error: %s | %s | %s",
+                          exception_type,
+                          exception_value,
+                          traceback)
+        self.driver.close()

utils/dates.py

@@ -0,0 +1,7 @@
+import datetime
+
+DATE_FMT = "%Y-%m-%d"
+
+
+def parse_date(date_str: str) -> datetime.datetime:
+    return datetime.datetime.strptime(date_str, DATE_FMT)