tinygent-graph 0.1.1__tar.gz

tinygent_graph-0.1.1/.gitignore

@@ -0,0 +1,210 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# examples/knowledge-graph
+.neo4j/

tinygent_graph-0.1.1/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: tinygent-graph
+Version: 0.1.1
+Requires-Dist: neo4j>=6.0.3
+Requires-Dist: tinygent

tinygent_graph-0.1.1/pyproject.toml

@@ -0,0 +1,17 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "tinygent-graph"
+version = "0.1.1"
+dependencies = [
+    "neo4j>=6.0.3",
+    "tinygent",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/tiny_graph"]
+
+[tool.uv.sources]
+tinygent = { workspace = true }

tinygent_graph-0.1.1/src/tiny_graph/__init__.py

@@ -0,0 +1,5 @@
+from .edge import TinyEdge
+from .graph import TinyMultiLayerGraph
+from .node import TinyNode
+
+__all__ = ['TinyEdge', 'TinyNode', 'TinyMultiLayerGraph']

tinygent_graph-0.1.1/src/tiny_graph/driver/__init__.py

@@ -0,0 +1,5 @@
+from .neo4j import Neo4jDriver
+
+__all__ = [
+    'Neo4jDriver',
+]

tinygent_graph-0.1.1/src/tiny_graph/driver/base.py

@@ -0,0 +1,21 @@
+from abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+from typing_extensions import LiteralString
+
+from tiny_graph.types.provider import GraphProvider
+
+
+class BaseDriver(ABC):
+    """Abstract base class for graph db drivers."""
+
+    provider: GraphProvider
+
+    @abstractmethod
+    async def execute_query(self, query: str | LiteralString, **kwargs: Any) -> Any:
+        raise NotImplementedError('Subclasses must implement this method.')
+
+    @abstractmethod
+    async def close(self) -> None:
+        raise NotImplementedError('Subclasses must implement this method.')

tinygent_graph-0.1.1/src/tiny_graph/driver/neo4j.py

@@ -0,0 +1,60 @@
+import logging
+from typing import Any
+from typing import cast
+
+from neo4j import AsyncGraphDatabase
+from neo4j import EagerResult
+from typing_extensions import LiteralString
+
+from tiny_graph.driver.base import BaseDriver
+from tiny_graph.types.provider import GraphProvider
+
+logger = logging.getLogger(__name__)
+
+
+class Neo4jDriver(BaseDriver):
+    provider = GraphProvider.NEO4J
+
+    def __init__(
+        self,
+        uri: str,
+        user: str,
+        password: str,
+    ) -> None:
+        self.uri = uri
+        self.user = user
+        self.password = password
+
+        self.__client = AsyncGraphDatabase.driver(
+            uri,
+            auth=(user, password),
+        )
+
+    async def health_check(self) -> None:
+        try:
+            await self.__client.verify_connectivity()
+            logger.debug('Neo4j connection is healthy')
+        except Exception as e:
+            logger.error('Neo4j health check failed: %s', e)
+            raise e
+
+    async def execute_query(
+        self, query: str | LiteralString, **kwargs: Any
+    ) -> EagerResult:
+        params = kwargs.pop('params', {})
+
+        if isinstance(query, str):
+            query = cast(LiteralString, query)
+
+        try:
+            result = await self.__client.execute_query(
+                query, parameters_=params, **kwargs
+            )
+        except Exception as e:
+            logger.error('Neo4j failed to execute query: %s with error: %s', query, e)
+            raise e
+        return result
+
+    async def close(self) -> None:
+        await self.__client.close()
+        logger.debug('Neo4j connection closed')

tinygent_graph-0.1.1/src/tiny_graph/edge.py

@@ -0,0 +1,40 @@
+from abc import ABC
+from abc import abstractmethod
+from datetime import datetime
+from typing import Any
+
+from pydantic import Field
+
+from tiny_graph.driver.base import BaseDriver
+from tiny_graph.helper import generate_uuid
+from tiny_graph.helper import get_current_timestamp
+from tinygent.core.types.base import TinyModel
+
+
+class TinyEdge(TinyModel, ABC):
+    uuid: str = Field(
+        description='unique edge identifier', default_factory=generate_uuid
+    )
+
+    name: str = Field(description='name of the edge, relation name')
+
+    subgraph_id: str = Field(..., description='subgraph identifier')
+
+    source_node_uuid: str
+
+    target_node_uuid: str
+
+    created_at: datetime = Field(default_factory=get_current_timestamp)
+
+    attributes: dict[str, Any] = Field(
+        default_factory=dict, description='Additional attributes of the node.'
+    )
+
+    @classmethod
+    @abstractmethod
+    def from_record(cls, record: dict) -> Any:
+        raise NotImplementedError('Subclasses must implement this method.')
+
+    @abstractmethod
+    async def save(self, driver: BaseDriver) -> Any:
+        raise NotImplementedError('Subclasses must implement this method.')

tinygent_graph-0.1.1/src/tiny_graph/graph/__init__.py

@@ -0,0 +1,5 @@
+from .multi_layer_graph.multi_layer_graph import TinyMultiLayerGraph
+
+__all__ = [
+    'TinyMultiLayerGraph',
+]

tinygent_graph-0.1.1/src/tiny_graph/graph/base.py

@@ -0,0 +1,37 @@
+from abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+from tiny_graph.driver.base import BaseDriver
+from tinygent.core.datamodels.embedder import AbstractEmbedder
+from tinygent.core.datamodels.llm import AbstractLLM
+from tinygent.core.datamodels.messages import BaseMessage
+
+
+class BaseGraph(ABC):
+    def __init__(
+        self,
+        llm: AbstractLLM,
+        embedder: AbstractEmbedder,
+        driver: BaseDriver,
+    ) -> None:
+        self.llm = llm
+        self.embedder = embedder
+        self.driver = driver
+
+    @abstractmethod
+    async def add_record(
+        self,
+        name: str,
+        data: str | dict | BaseMessage,
+        description: str,
+        *,
+        uuid: str | None = None,
+        subgraph_id: str | None = None,
+        **kwargs,
+    ) -> Any:
+        pass
+
+    @abstractmethod
+    async def close(self) -> None:
+        pass

tinygent_graph-0.1.1/src/tiny_graph/graph/multi_layer_graph/__init__.py

@@ -0,0 +1,5 @@
+from .search.search import search
+
+__all__ = [
+    'search',
+]

tinygent_graph-0.1.1/src/tiny_graph/graph/multi_layer_graph/core/edge.py

@@ -0,0 +1,14 @@
+from tiny_graph.graph.multi_layer_graph.edges import TinyEntityEdge
+from tinygent.core.datamodels.embedder import AbstractEmbedder
+
+
+async def entity_edge_batch_embeddings(
+    embedder: AbstractEmbedder, edges: list[TinyEntityEdge]
+) -> list[TinyEntityEdge]:
+    if not edges:
+        return []
+
+    embeddings = await embedder.aembed_batch([e.fact for e in edges])
+    for edge, emb in zip(edges, embeddings, strict=True):
+        edge.fact_embedding = emb
+    return edges

tinygent_graph-0.1.1/src/tiny_graph/graph/multi_layer_graph/core/node.py

@@ -0,0 +1,27 @@
+from tiny_graph.graph.multi_layer_graph.nodes import TinyClusterNode
+from tiny_graph.graph.multi_layer_graph.nodes import TinyEntityNode
+from tinygent.core.datamodels.embedder import AbstractEmbedder
+
+
+async def cluster_node_batch_embeddings(
+    embedder: AbstractEmbedder, clusters: list[TinyClusterNode]
+) -> list[TinyClusterNode]:
+    if not clusters:
+        return []
+
+    embeddings = await embedder.aembed_batch([c.name for c in clusters])
+    for cluster, emb in zip(clusters, embeddings, strict=True):
+        cluster.name_embedding = emb
+    return clusters
+
+
+async def entity_node_batch_embeddings(
+    embedder: AbstractEmbedder, entities: list[TinyEntityNode]
+) -> list[TinyEntityNode]:
+    if not entities:
+        return []
+
+    embeddings = await embedder.aembed_batch([e.name for e in entities])
+    for entity, emb in zip(entities, embeddings, strict=True):
+        entity.name_embedding = emb
+    return entities

tinygent_graph-0.1.1/src/tiny_graph/graph/multi_layer_graph/datamodels/clients.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+
+from tiny_graph.driver.base import BaseDriver
+from tinygent.core.datamodels.cross_encoder import AbstractCrossEncoder
+from tinygent.core.datamodels.embedder import AbstractEmbedder
+from tinygent.core.datamodels.llm import AbstractLLM
+
+
+@dataclass
+class TinyGraphClients:
+    driver: BaseDriver
+    llm: AbstractLLM
+    embedder: AbstractEmbedder
+    cross_encoder: AbstractCrossEncoder
+
+    @property
+    def safe_embed_model(self) -> str:
+        return self.embedder.model.replace('-', '_')

tinygent_graph-0.1.1/src/tiny_graph/graph/multi_layer_graph/datamodels/extract_nodes.py

@@ -0,0 +1,23 @@
+from pydantic import Field
+
+from tinygent.core.types.base import TinyModel
+
+
+class ExtractedEntity(TinyModel):
+    name: str = Field(..., description='Name of the extracted entity')
+    entity_type_id: int = Field(
+        description='ID of the classified entity type. '
+        'Must be one of the provided entity_type_id integers.',
+    )
+
+
+class ExtractedEntities(TinyModel):
+    extracted_entities: list[ExtractedEntity] = Field(
+        ..., description='List of extracted entities'
+    )
+
+
+class MissedEntities(TinyModel):
+    missed_entities: list[str] = Field(
+        ..., description="Names of entities that weren't extracted"
+    )

tinygent_graph-0.1.1/src/tiny_graph/graph/multi_layer_graph/edges.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from pydantic import Field
+
+from tiny_graph.driver.base import BaseDriver
+from tiny_graph.edge import TinyEdge
+from tiny_graph.graph.multi_layer_graph.types import EdgeType
+from tiny_graph.graph.multi_layer_graph.utils.model_repr import compact_model_repr
+from tiny_graph.helper import parse_db_date
+from tinygent.core.datamodels.embedder import AbstractEmbedder
+from tinygent.utils.yaml import json
+
+
+class TinyEntityEdge(TinyEdge):
+    fact: str = Field(
+        description='fact representing the edge and nodes that it connects'
+    )
+
+    fact_embedding: list[float] | None = Field(
+        default=None, description='embedding of the fact'
+    )
+
+    events: list[str] = Field(
+        default=[],
+        description='list of episode ids that reference these entity edges',
+    )
+
+    expired_at: datetime | None = Field(
+        default=None, description='datetime of when the node was invalidated'
+    )
+
+    valid_at: datetime | None = Field(
+        default=None, description='datetime of when the fact became true'
+    )
+
+    invalid_at: datetime | None = Field(
+        default=None, description='datetime of when the fact stopped being true'
+    )
+
+    async def save(self, driver: BaseDriver) -> str:
+        from tiny_graph.graph.multi_layer_graph.queries.edge_queries import (
+            create_entity_edge,
+        )
+
+        args = {
+            'edge_uuid': self.uuid,
+            'subgraph_id': self.subgraph_id,
+            'source_node_uuid': self.source_node_uuid,
+            'target_node_uuid': self.target_node_uuid,
+            'created_at': self.created_at,
+            'name': self.name,
+            'fact': self.fact,
+            'fact_embedding': self.fact_embedding,
+            'events': self.events,
+            'expired_at': self.expired_at,
+            'valid_at': self.valid_at,
+            'invalid_at': self.invalid_at,
+            'attributes': json.dumps(self.attributes),
+        }
+
+        return await driver.execute_query(
+            query=create_entity_edge(driver.provider),
+            **args,
+        )
+
+    async def embed(self, embedder: AbstractEmbedder) -> list[float]:
+        self.fact_embedding = await embedder.aembed(self.fact)
+        return self.fact_embedding
+
+    @classmethod
+    def from_record(cls, record: dict) -> TinyEntityEdge:
+        return TinyEntityEdge(
+            uuid=record['uuid'],
+            subgraph_id=record['subgraph_id'],
+            source_node_uuid=record['source_node_uuid'],
+            target_node_uuid=record['target_node_uuid'],
+            created_at=parse_db_date(record['created_at']),
+            name=record['name'],
+            fact=record['fact'],
+            fact_embedding=record.get('fact_embedding'),
+            events=record.get('events', []),
+            expired_at=(
+                parse_db_date(record['expired_at']) if record.get('expired_at') else None
+            ),
+            valid_at=(
+                parse_db_date(record['valid_at']) if record.get('valid_at') else None
+            ),
+            invalid_at=(
+                parse_db_date(record['invalid_at']) if record.get('invalid_at') else None
+            ),
+            attributes=(
+                json.loads(record['attributes'])
+                if isinstance(record.get('attributes'), str)
+                else record.get('attributes', {})
+            ),
+        )
+
+    @classmethod
+    async def find_by_targets(
+        cls, driver: BaseDriver, source_uuid: str, target_uuid: str
+    ) -> list[TinyEntityEdge]:
+        from tiny_graph.graph.multi_layer_graph.queries.edge_queries import (
+            find_entity_edge_by_targets,
+        )
+
+        query = find_entity_edge_by_targets(driver.provider)
+
+        results, _, _ = await driver.execute_query(
+            query,
+            source_uuid=source_uuid,
+            target_uuid=target_uuid,
+        )
+        return [TinyEntityEdge.from_record(r) for r in results]
+
+    def __repr__(self) -> str:
+        return compact_model_repr(self)
+
+
+class TinyClusterEdge(TinyEdge):
+    name: str = Field(default=EdgeType.HAS_MEMBER.value, frozen=True)
+
+    async def save(self, driver: BaseDriver) -> str:
+        from tiny_graph.graph.multi_layer_graph.queries.edge_queries import (
+            create_cluster_edge,
+        )
+
+        args = {
+            'uuid': self.uuid,
+            'subgraph_id': self.subgraph_id,
+            'created_at': self.created_at,
+            'cluster_node_uuid': self.source_node_uuid,
+            'entity_node_uuid': self.target_node_uuid,
+        }
+
+        return await driver.execute_query(
+            query=create_cluster_edge(driver.provider),
+            **args,
+        )
+
+    @classmethod
+    def from_record(cls, record: dict) -> TinyClusterEdge:
+        return TinyClusterEdge(
+            uuid=record['uuid'],
+            subgraph_id=record['subgraph_id'],
+            source_node_uuid=record['source_node_uuid'],
+            target_node_uuid=record['target_node_uuid'],
+            created_at=parse_db_date(record['created_at']),
+            name=record['name'],
+        )
+
+
+class TinyEventEdge(TinyEdge):
+    name: str = Field(default=EdgeType.MENTIONS.value, frozen=True)
+
+    async def save(self, driver: BaseDriver) -> str:
+        from tiny_graph.graph.multi_layer_graph.queries.edge_queries import (
+            create_event_edge,
+        )
+
+        args = {
+            'uuid': self.uuid,
+            'subgraph_id': self.subgraph_id,
+            'created_at': self.created_at,
+            'event_node_uuid': self.source_node_uuid,
+            'entity_node_uuid': self.target_node_uuid,
+        }
+
+        return await driver.execute_query(
+            query=create_event_edge(driver.provider),
+            **args,
+        )
+
+    @classmethod
+    def from_record(cls, record: dict) -> TinyClusterEdge:
+        return TinyClusterEdge(
+            uuid=record['uuid'],
+            subgraph_id=record['subgraph_id'],
+            source_node_uuid=record['source_node_uuid'],
+            target_node_uuid=record['target_node_uuid'],
+            created_at=parse_db_date(record['created_at']),
+            name=record['name'],
+        )