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

vedana_core/app.py

@@ -0,0 +1,78 @@
+from dataclasses import dataclass
+
+import sqlalchemy.ext.asyncio as sa_aio
+from async_lru import alru_cache
+from jims_core.app import JimsApp
+from loguru import logger
+
+from vedana_core.data_model import DataModel
+from vedana_core.db import get_sessionmaker
+from vedana_core.graph import Graph, MemgraphGraph
+from vedana_core.vts import VectorStore, PGVectorStore
+from vedana_core.rag_pipeline import RagPipeline
+from vedana_core.start_pipeline import StartPipeline
+from vedana_core.settings import settings as core_settings
+
+
+@dataclass
+class VedanaApp:
+    sessionmaker: sa_aio.async_sessionmaker[sa_aio.AsyncSession]
+
+    graph: Graph
+    vts: VectorStore
+    data_model: DataModel
+    pipeline: RagPipeline
+    start_pipeline: StartPipeline
+
+
+@alru_cache
+async def make_vedana_app() -> VedanaApp:
+    sessionmaker = get_sessionmaker()
+
+    graph = MemgraphGraph(
+        core_settings.memgraph_uri,
+        core_settings.memgraph_user,
+        core_settings.memgraph_pwd,
+    )
+
+    vts = PGVectorStore(
+        sessionmaker=sessionmaker,
+    )
+
+    data_model = DataModel(sessionmaker=sessionmaker)
+
+    pipeline = RagPipeline(
+        graph=graph,
+        vts=vts,
+        data_model=data_model,
+        logger=logger,
+        threshold=0.8,
+    )
+
+    start_pipeline = StartPipeline(data_model=data_model)
+
+    return VedanaApp(
+        sessionmaker=sessionmaker,
+        graph=graph,
+        vts=vts,
+        data_model=data_model,
+        pipeline=pipeline,
+        start_pipeline=start_pipeline,
+    )
+
+
+@alru_cache
+async def make_jims_app() -> JimsApp:
+    vedana_app = await make_vedana_app()
+
+    app = JimsApp(
+        sessionmaker=vedana_app.sessionmaker,
+        pipeline=vedana_app.pipeline,
+        conversation_start_pipeline=vedana_app.start_pipeline,
+    )
+
+    return app
+
+
+# This creates a async coroutine which will be evaluated in the event loop
+app = make_jims_app()

vedana_core/data_model.py

@@ -0,0 +1,465 @@
+import logging
+from dataclasses import dataclass
+
+import sqlalchemy.ext.asyncio as sa_aio
+from sqlalchemy import select
+from vedana_etl.catalog import (
+    dm_anchor_attributes,
+    dm_anchors,
+    dm_conversation_lifecycle,
+    dm_link_attributes,
+    dm_links,
+    dm_prompts,
+    dm_queries,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Attribute:
+    name: str
+    description: str
+    example: str
+    dtype: str
+    query: str
+    embeddable: bool = False
+    embed_threshold: float = 0
+
+
+@dataclass
+class Anchor:
+    noun: str
+    description: str
+    id_example: str
+    query: str
+    attributes: list[Attribute]
+
+    def __str__(self) -> str:
+        return self.noun
+
+
+@dataclass
+class Link:
+    anchor_from: Anchor
+    anchor_to: Anchor
+    sentence: str
+    description: str
+    query: str
+    attributes: list[Attribute]
+    has_direction: bool = False
+    anchor_from_link_attr_name: str = ""
+    anchor_to_link_attr_name: str = ""
+
+
+@dataclass
+class Query:
+    name: str
+    example: str
+
+
+@dataclass
+class ConversationLifecycleEvent:
+    event: str
+    text: str
+
+
+@dataclass
+class Prompt:
+    name: str
+    text: str
+
+
+class DataModel:
+    """
+    DataModel, read from SQL tables at runtime
+    """
+
+    def __init__(self, sessionmaker: sa_aio.async_sessionmaker[sa_aio.AsyncSession]) -> None:
+        self.sessionmaker = sessionmaker
+
+    @classmethod
+    def create(cls, sessionmaker) -> "DataModel":
+        return cls(sessionmaker=sessionmaker)
+
+    async def get_anchors(self) -> list[Anchor]:
+        """Read anchors from dm_anchors table."""
+        # .data_table is TableStoreDB's attribute
+        anchors_table = dm_anchors.store.data_table  # type: ignore[attr-defined]
+        anchors_attr_table = dm_anchor_attributes.store.data_table  # type: ignore[attr-defined]
+
+        async with self.sessionmaker() as session:
+            join_query = select(
+                anchors_table.c.noun,
+                anchors_table.c.description.label("anchor_description"),
+                anchors_table.c.id_example,
+                anchors_table.c.query.label("anchor_query"),
+                anchors_attr_table.c.attribute_name,
+                anchors_attr_table.c.description.label("attr_description"),
+                anchors_attr_table.c.data_example,
+                anchors_attr_table.c.embeddable,
+                anchors_attr_table.c.query.label("attr_query"),
+                anchors_attr_table.c.dtype,
+                anchors_attr_table.c.embed_threshold,
+            ).select_from(
+                anchors_table.join(  # left join
+                    anchors_attr_table,
+                    anchors_table.c.noun == anchors_attr_table.c.anchor,
+                    isouter=True,
+                )
+            )
+            result = (await session.execute(join_query)).fetchall()
+
+            anchors = {}
+            for row in result:
+                noun = row.noun
+                if noun not in anchors:
+                    anchors[noun] = Anchor(
+                        noun=noun,
+                        description=row.anchor_description,
+                        id_example=row.id_example,
+                        query=row.anchor_query,
+                        attributes=[],
+                    )
+
+                # Add attribute if it exists (attribute_name will be None for anchors without attributes)
+                if row.attribute_name is not None:
+                    anchors[noun].attributes.append(
+                        Attribute(
+                            name=row.attribute_name,
+                            description=row.attr_description if row.attr_description else "",
+                            example=row.data_example if row.data_example else "",
+                            embeddable=row.embeddable if row.embeddable is not None else False,
+                            query=row.attr_query if row.attr_query else "",
+                            dtype=row.dtype if row.dtype else "",
+                            embed_threshold=row.embed_threshold if row.embed_threshold is not None else 1.0,
+                        )
+                    )
+
+            return list(anchors.values())
+
+    async def get_links(self, anchors_dict: dict[str, Anchor] | None = None) -> list[Link]:
+        """Read links from dm_links table."""
+        links_table = dm_links.store.data_table  # type: ignore[attr-defined]
+        links_attr_table = dm_link_attributes.store.data_table  # type: ignore[attr-defined]
+
+        if anchors_dict is None:
+            anchors = await self.get_anchors()
+            anchors_dict = {anchor.noun: anchor for anchor in anchors}
+
+        async with self.sessionmaker() as session:
+            join_query = select(
+                links_table.c.anchor1,
+                links_table.c.anchor2,
+                links_table.c.sentence,
+                links_table.c.description.label("link_description"),
+                links_table.c.query.label("link_query"),
+                links_table.c.anchor1_link_column_name,
+                links_table.c.anchor2_link_column_name,
+                links_table.c.has_direction,
+                links_attr_table.c.attribute_name,
+                links_attr_table.c.description.label("attr_description"),
+                links_attr_table.c.data_example,
+                links_attr_table.c.embeddable,
+                links_attr_table.c.query.label("attr_query"),
+                links_attr_table.c.dtype,
+                links_attr_table.c.embed_threshold,
+            ).select_from(
+                links_table.join(  # left join
+                    links_attr_table,
+                    links_table.c.sentence == links_attr_table.c.link,
+                    isouter=True,
+                )
+            )
+
+            result = (await session.execute(join_query)).fetchall()
+
+            links = {}
+            for row in result:
+                sentence = row.sentence
+                if sentence not in links:
+                    anchor_from = anchors_dict.get(row.anchor1)
+                    anchor_to = anchors_dict.get(row.anchor2)
+                    if anchor_from is None or anchor_to is None:
+                        logger.error(f'Link {sentence} has invalid connection "{row.anchor1} - {row.anchor2}"')
+                        continue
+
+                    links[sentence] = Link(
+                        anchor_from=anchor_from,
+                        anchor_to=anchor_to,
+                        anchor_from_link_attr_name=row.anchor1_link_column_name,
+                        anchor_to_link_attr_name=row.anchor2_link_column_name,
+                        sentence=sentence,
+                        description=row.link_description,
+                        query=row.link_query,
+                        has_direction=bool(row.has_direction) if row.has_direction is not None else False,
+                        attributes=[],
+                    )
+
+                # Add attribute if it exists (attribute_name will be None for anchors without attributes)
+                if row.attribute_name is not None:
+                    links[sentence].attributes.append(
+                        Attribute(
+                            name=row.attribute_name,
+                            description=row.attr_description if row.attr_description else "",
+                            example=row.data_example if row.data_example else "",
+                            embeddable=row.embeddable if row.embeddable is not None else False,
+                            query=row.attr_query if row.attr_query else "",
+                            dtype=row.dtype if row.dtype else "",
+                            embed_threshold=row.embed_threshold if row.embed_threshold is not None else 1.0,
+                        )
+                    )
+
+            return list(links.values())
+
+    async def get_queries(self) -> list[Query]:
+        try:
+            queries_table = dm_queries.store.data_table  # type: ignore[attr-defined]
+            async with self.sessionmaker() as session:
+                result = (await session.execute(select(queries_table))).fetchall()
+                return [Query(name=row.query_name, example=row.query_example) for row in result]
+        except Exception:
+            return []
+
+    async def get_conversation_lifecycle_events(self) -> list[ConversationLifecycleEvent]:
+        try:
+            lifecycle_table = dm_conversation_lifecycle.store.data_table  # type: ignore[attr-defined]
+            async with self.sessionmaker() as session:
+                result = (await session.execute(select(lifecycle_table))).fetchall()
+                return [ConversationLifecycleEvent(event=row.event, text=row.text) for row in result]
+        except Exception:
+            return []
+
+    async def conversation_lifecycle_events(self) -> dict[str, str]:
+        cl = await self.get_conversation_lifecycle_events()
+        return {c.event: c.text for c in cl}
+
+    async def get_prompts(self) -> list[Prompt]:
+        try:
+            prompts_table = dm_prompts.store.data_table  # type: ignore[attr-defined]
+            async with self.sessionmaker() as session:
+                result = (await session.execute(select(prompts_table))).fetchall()
+                return [Prompt(name=row.name, text=row.text) for row in result]
+        except Exception:
+            return []
+
+    async def prompt_templates(self) -> dict[str, str]:
+        prompts = await self.get_prompts()
+        return {p.name: p.text for p in prompts}
+
+    async def vector_indices(self) -> list[tuple[str, str, str, float]]:
+        """
+        returns list
+        ("anchor", anchor.noun, anchor.attribute, anchor.th) +
+        ("edge", link.sentence, link.attribute, link.th)
+        for all embeddable attributes
+        """
+        anchors = await self.get_anchors()
+        links = await self.get_links(anchors_dict={a.noun: a for a in anchors})
+
+        a_i = [
+            ("anchor", anchor.noun, attr.name, attr.embed_threshold)
+            for anchor in anchors
+            for attr in anchor.attributes
+            if attr.embeddable
+        ]
+        l_i = [
+            ("edge", link.sentence, attr.name, attr.embed_threshold)
+            for link in links
+            for attr in link.attributes
+            if attr.embeddable
+        ]
+        return a_i + l_i
+
+    async def anchor_links(self, anchor_noun: str) -> list[Link]:
+        """all links that connect to/from this anchor"""
+        links = await self.get_links()
+        return [
+            link
+            for link in links
+            if (link.anchor_from.noun == anchor_noun and link.anchor_from_link_attr_name)
+            or (link.anchor_to.noun == anchor_noun and link.anchor_to_link_attr_name)
+        ]
+
+    async def to_text_descr(
+        self,
+        anchor_nouns: list[str] | None = None,
+        link_sentences: list[str] | None = None,
+        anchor_attribute_names: list[str] | None = None,
+        link_attribute_names: list[str] | None = None,
+        query_names: list[str] | None = None,
+    ) -> str:
+        """Create a text description of the data model, optionally filtered.
+
+        Args:
+            anchor_nouns: List of anchor nouns to include. If None, includes all.
+            link_sentences: List of link sentences to include. If None, includes all.
+            anchor_attribute_names: List of anchor attribute names to include. If None, includes all.
+            link_attribute_names: List of link attribute names to include. If None, includes all.
+            query_names: List of query names to include. If None, includes all.
+
+        Returns:
+            A formatted string description of the data model.
+        """
+        anchors = await self.get_anchors()
+        links = await self.get_links(anchors_dict={a.noun: a for a in anchors})
+        queries = await self.get_queries()
+        dm_templates = await self.prompt_templates()
+
+        # Convert to sets for efficient lookup, None means include all
+        anchor_set = set(anchor_nouns) if anchor_nouns is not None else None
+        link_set = set(link_sentences) if link_sentences is not None else None
+        anchor_attr_set = set(anchor_attribute_names) if anchor_attribute_names is not None else None
+        link_attr_set = set(link_attribute_names) if link_attribute_names is not None else None
+        query_set = set(query_names) if query_names is not None else None
+
+        # Filter anchors
+        filtered_anchors = [anchor for anchor in anchors if anchor_set is None or anchor.noun in anchor_set]
+
+        # Create a map for quick anchor lookup (for link filtering)
+        anchors_map = {a.noun: a for a in filtered_anchors}
+
+        # Filter links (only include if both anchors are in filtered set)
+        filtered_links = [
+            link
+            for link in links
+            if (link_set is None or link.sentence in link_set)
+            and link.anchor_from.noun in anchors_map
+            and link.anchor_to.noun in anchors_map
+        ]
+
+        anchor_descr = "\n".join(
+            dm_templates.get("dm_anchor_descr_template", dm_anchor_descr_template).format(anchor=anchor)
+            for anchor in filtered_anchors
+        )
+
+        anchor_attrs_descr = "\n".join(
+            dm_templates.get("dm_attr_descr_template", dm_attr_descr_template).format(anchor=anchor, attr=attr)
+            for anchor in filtered_anchors
+            for attr in anchor.attributes
+            if anchor_attr_set is None or attr.name in anchor_attr_set
+        )
+
+        link_descr = "\n".join(
+            dm_templates.get("dm_link_descr_template", dm_link_descr_template).format(link=link)
+            for link in filtered_links
+        )
+
+        link_attrs_descr = "\n".join(
+            dm_templates.get("dm_link_attr_descr_template", dm_link_attr_descr_template).format(link=link, attr=attr)
+            for link in filtered_links
+            for attr in link.attributes
+            if link_attr_set is None or attr.name in link_attr_set
+        )
+
+        filtered_queries = [q for q in queries if query_set is None or q.name in query_set]
+        queries_descr = "\n".join(
+            dm_templates.get("dm_query_descr_template", dm_query_descr_template).format(query=query)
+            for query in filtered_queries
+        )
+
+        dm_template = dm_templates.get("dm_descr_template", dm_descr_template)
+
+        return dm_template.format(
+            anchors=anchor_descr,
+            anchor_attrs=anchor_attrs_descr,
+            links=link_descr,
+            link_attrs=link_attrs_descr,
+            queries=queries_descr,
+        )
+
+    async def to_compact_json(self) -> dict:
+        anchors = await self.get_anchors()
+        links = await self.get_links(anchors_dict={a.noun: a for a in anchors})
+        queries = await self.get_queries()
+
+        descr = {
+            "anchors": [
+                {
+                    "name": a.noun,
+                    "description": a.description,
+                    "example": a.id_example,
+                    "attributes": [
+                        {
+                            "attr_name": aa.name,
+                            "attr_description": aa.description,
+                        }
+                        for aa in a.attributes
+                    ],
+                }
+                for a in anchors
+            ],
+            "links": [
+                {
+                    "from": li.anchor_from,
+                    "to": li.anchor_to,
+                    "sentence": li.sentence,
+                    "description": li.description,
+                    "attributes": [
+                        {
+                            "attr_name": la.name,
+                            "attr_description": la.description,
+                        }
+                        for la in li.attributes
+                    ],
+                }
+                for li in links
+            ],
+            "queries": {i: q.name for i, q in enumerate(queries, start=1)},
+        }
+        return descr
+
+
+# default templates
+dm_descr_template = """\
+## Узлы:
+{anchors}
+
+## Атрибуты узлов:
+{anchor_attrs}
+
+## Связи между узлами:
+{links}
+
+## Атрибуты связей:
+{link_attrs}
+
+## Типичные вопросы:
+{queries}
+"""
+
+dm_anchor_descr_template = (
+    "- {anchor.noun}: {anchor.description}; пример ID: {anchor.id_example}; запрос для получения: {anchor.query}"
+)
+dm_attr_descr_template = (
+    "- {anchor.noun}.{attr.name}: {attr.description}; пример: {attr.example}; запрос для получения: {attr.query}"
+)
+dm_link_descr_template = "- {link.sentence}: {link.description}; пример запроса: {link.query}"
+dm_link_attr_descr_template = (
+    "- {link.sentence}.{attr.name}: {attr.description}; пример: {attr.example}; запрос для получения: {attr.query}"
+)
+dm_query_descr_template = "- {query.name}\n{query.example}"
+
+# Compact templates (without cypher queries)
+dm_compact_descr_template = """\
+## Узлы:
+{anchors}
+
+## Атрибуты узлов:
+{anchor_attrs}
+
+## Связи между узлами:
+{links}
+
+## Атрибуты связей:
+{link_attrs}
+
+## Сценарии вопросов:
+{queries}
+"""
+
+dm_compact_anchor_descr_template = "- {anchor.noun}: {anchor.description}"
+dm_compact_attr_descr_template = "- {anchor.noun}.{attr.name}: {attr.description}"
+dm_compact_link_descr_template = "- {link.sentence}: {link.description}"
+dm_compact_link_attr_descr_template = "- {link.sentence}.{attr.name}: {attr.description}"
+dm_compact_query_descr_template = "- {query.name}"