vedana-core 0.1.0.dev3__py3-none-any.whl

vedana_core/vts.py

@@ -0,0 +1,167 @@
+import abc
+import logging
+from typing import Sequence
+
+import numpy as np
+from neo4j import AsyncGraphDatabase, Record, RoutingControl
+from opentelemetry import trace
+from sqlalchemy import RowMapping, select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from vedana_etl.catalog import edges, nodes, rag_anchor_embeddings, rag_edge_embeddings
+
+from vedana_core.db import get_sessionmaker
+
+logger = logging.getLogger(__name__)
+tracer = trace.get_tracer(__name__)
+
+
+class VectorStore(abc.ABC):
+    async def vector_search(
+        self,
+        label: str,
+        prop_type: str,
+        prop_name: str,
+        embedding: list[float],
+        threshold: float,
+        top_n: int = 5,
+    ):
+        raise NotImplementedError
+
+
+class MemgraphVectorStore(VectorStore):
+    """
+    Use Memgraph vector_search capabilities. Requires vector indices to be created separately.
+    """
+
+    def __init__(self, uri: str, user: str, pwd: str, db_name: str = "") -> None:
+        self.driver = AsyncGraphDatabase.driver(uri, auth=(user, pwd), database=db_name)
+        # await self.driver.verify_connectivity()
+        self.driver_uri = uri
+        self.auth = (user, pwd)
+
+    async def vector_search(
+        self,
+        label: str,
+        prop_type: str,
+        prop_name: str,
+        embedding: np.ndarray | list[float],
+        threshold: float,
+        top_n: int = 5,
+    ) -> list[Record]:
+        with tracer.start_as_current_span("memgraph.vector_search") as span:
+            span.set_attribute("memgraph.label", label)
+            span.set_attribute("memgraph.prop_type", prop_type)
+            span.set_attribute("memgraph.prop_name", prop_name)
+            span.set_attribute("memgraph.top_n", top_n)
+            span.set_attribute("memgraph.threshold", threshold)
+
+            if prop_type == "edge":
+                query = (
+                    "CALL vector_search.search_edges($idx_name, $top_n, $embedding) "
+                    "YIELD similarity, edge "
+                    "WITH similarity, edge "
+                    "WHERE similarity > $threshold "
+                    "RETURN similarity, edge, startNode(edge) AS start, endNode(edge) AS end;"
+                )
+            else:  # node
+                query = (
+                    "CALL vector_search.search($idx_name, $top_n, $embedding) "
+                    "YIELD similarity, node "
+                    "WITH similarity, node WHERE similarity > $threshold RETURN *"
+                )
+
+            span.set_attribute("memgraph.query", query)
+
+            idx_name = f"{label}_{prop_name}_embed_idx".replace(" ", "_")
+            res = await self.driver.execute_query(
+                query,
+                idx_name=idx_name,
+                top_n=top_n,
+                embedding=embedding,
+                threshold=threshold,
+                routing_=RoutingControl.READ,
+            )
+            return res.records
+
+
+class PGVectorStore(VectorStore):
+    def __init__(
+        self,
+        sessionmaker: async_sessionmaker[AsyncSession] | None = None,
+    ) -> None:
+        self._sessionmaker: async_sessionmaker[AsyncSession] = sessionmaker or get_sessionmaker()
+        self.rag_anchor_embeddings_table = rag_anchor_embeddings.store.data_table  # type: ignore[attr-defined]
+        self.rag_edge_embeddings_table = rag_edge_embeddings.store.data_table  # type: ignore[attr-defined]
+        self.node_table = nodes.store.data_table  # type: ignore[attr-defined]
+        self.edge_table = edges.store.data_table  # type: ignore[attr-defined]
+
+    async def vector_search(
+        self,
+        label: str,
+        prop_type: str,
+        prop_name: str,
+        embedding: np.ndarray | list[float],
+        threshold: float,
+        top_n: int = 5,
+    ) -> Sequence[RowMapping]:
+        with tracer.start_as_current_span("pgvector.vector_search") as span:
+            span.set_attribute("pgvector.label", label)
+            span.set_attribute("pgvector.prop_type", prop_type)
+            span.set_attribute("pgvector.prop_name", prop_name)
+            span.set_attribute("pgvector.top_n", top_n)
+            span.set_attribute("pgvector.threshold", threshold)
+
+            async with self._sessionmaker() as session:
+                if prop_type == "edge":
+                    distance = self.rag_edge_embeddings_table.c.embedding.cosine_distance(embedding)
+                    similarity = (1 - distance).label("similarity")
+
+                    stmt = (
+                        select(
+                            similarity,
+                            self.edge_table.c.from_node_id,
+                            self.edge_table.c.to_node_id,
+                            self.edge_table.c.edge_label,
+                            self.edge_table.c.attributes.label("edge"),
+                        )
+                        .select_from(
+                            self.rag_edge_embeddings_table.join(
+                                self.edge_table,
+                                (self.rag_edge_embeddings_table.c.from_node_id == self.edge_table.c.from_node_id)
+                                & (self.rag_edge_embeddings_table.c.to_node_id == self.edge_table.c.to_node_id)
+                                & (self.rag_edge_embeddings_table.c.edge_label == self.edge_table.c.edge_label),
+                            )
+                        )
+                        .where(self.rag_edge_embeddings_table.c.edge_label == label)
+                        .where(self.rag_edge_embeddings_table.c.attribute_name == prop_name)
+                        .where(similarity > threshold)
+                        .order_by(distance)
+                        .limit(top_n)
+                    )
+
+                else:  # node
+                    distance = self.rag_anchor_embeddings_table.c.embedding.cosine_distance(embedding)
+                    similarity = (1 - distance).label("similarity")
+
+                    stmt = (
+                        select(
+                            similarity,
+                            self.node_table.c.node_id,
+                            self.node_table.c.node_type,
+                            self.node_table.c.attributes.label("node"),
+                        )
+                        .select_from(
+                            self.rag_anchor_embeddings_table.join(
+                                self.node_table, self.rag_anchor_embeddings_table.c.node_id == self.node_table.c.node_id
+                            )
+                        )
+                        .where(self.rag_anchor_embeddings_table.c.label == label)
+                        .where(self.rag_anchor_embeddings_table.c.attribute_name == prop_name)
+                        .where(similarity > threshold)
+                        .order_by(distance)
+                        .limit(top_n)
+                    )
+
+                span.set_attribute("pgvector.query", str(stmt))
+                res = await session.execute(stmt)
+                return res.mappings().all()

vedana_core-0.1.0.dev3.dist-info/METADATA

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: vedana-core
+Version: 0.1.0.dev3
+Summary: Semantic Graph RAG App
+Author-email: Andrey Tatarinov <a@tatarinov.co>, Timur Sheydaev <tsheyd@epoch8.co>
+Requires-Python: >=3.12
+Requires-Dist: aiohttp>=3.11.18
+Requires-Dist: aioitertools>=0.12.0
+Requires-Dist: aiosqlite>=0.21.0
+Requires-Dist: alembic>=1.13.0
+Requires-Dist: async-lru>=2.0.5
+Requires-Dist: asyncpg>=0.29.0
+Requires-Dist: grist-api>=0.1.1
+Requires-Dist: jims-backoffice
+Requires-Dist: jims-core
+Requires-Dist: jims-telegram
+Requires-Dist: jims-tui
+Requires-Dist: litellm>=1.79.0
+Requires-Dist: neo4j>=5.28.1
+Requires-Dist: openai>=2.8.0
+Requires-Dist: pandas>=2.2.3
+Requires-Dist: psycopg2-binary>=2.9.10
+Requires-Dist: pydantic-settings>=2.9.1
+Requires-Dist: sentry-sdk[fastapi,opentelemetry]>=2.32.0
+Requires-Dist: sqlalchemy>=2.0.41
+Requires-Dist: vedana-etl
+Description-Content-Type: text/markdown
+
+# Vedana

vedana_core-0.1.0.dev3.dist-info/RECORD

@@ -0,0 +1,17 @@
+vedana_core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vedana_core/app.py,sha256=LmolsEf6_INsDWpeSaYDxoZ9v029K-hmlRvUV_65VT0,1893
+vedana_core/data_model.py,sha256=O24Rzwh712fhj4al71S2UM9RUFbg1UV5fojvPooH6UI,17589
+vedana_core/data_provider.py,sha256=r3yVKKm7Bo6tN0ZnmlIyo-Ndd5KNhBRPaC0wb895Xl0,18711
+vedana_core/db.py,sha256=Wi71gXBrqPv6N6MIjaaHHFVmGIjr3LiRs5AeduDbGLc,1143
+vedana_core/graph.py,sha256=uYYRdfgX5rXI30cdNvXpbevDo0fEzRcTrwIu6h11Q2k,10800
+vedana_core/llm.py,sha256=oRPPoQPOLhf17tipekOMcFjKTHUhryshmvnpC9NA5lQ,8948
+vedana_core/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vedana_core/rag_agent.py,sha256=kYYgVe7tLghLaaeaLUPJU8hS-LEZ1knx7As3oH_Wfks,8321
+vedana_core/rag_pipeline.py,sha256=eOnWYph-2qZufw4LzS6SeKWf69I7IKX_ynTd5H4noNY,13750
+vedana_core/settings.py,sha256=X78c46zppz6kuFIJx4xGUoaerlomPYz8XrYqeIOMy8c,892
+vedana_core/start_pipeline.py,sha256=SJWBCbod88z1Qv8st3bPAsqgGK7cRSwxOJcoWGRXvLM,597
+vedana_core/utils.py,sha256=blNd44rHLfVjWvKJGHx5q6RJPHHGBHE_8DYRUdYUX24,1189
+vedana_core/vts.py,sha256=JzIhb2bw1Z4fJrY7nn3FewIzRm2zZZVOUD8kgcgiucU,6981
+vedana_core-0.1.0.dev3.dist-info/METADATA,sha256=ctotRu_F0biJqNDOFV1b2A4Z_MJmNM8mhGNgLk7tACI,899
+vedana_core-0.1.0.dev3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+vedana_core-0.1.0.dev3.dist-info/RECORD,,

vedana_core-0.1.0.dev3.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any