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

vedana_core/rag_agent.py

@@ -0,0 +1,234 @@
+import enum
+import json
+import logging
+import re
+from dataclasses import dataclass
+from itertools import islice
+from typing import Any, Mapping, Type
+
+import neo4j
+import neo4j.graph
+from jims_core.thread.thread_context import ThreadContext
+from pydantic import BaseModel, Field, create_model
+
+from vedana_core.graph import Graph, Record
+from vedana_core.vts import VectorStore
+from vedana_core.llm import LLM, Tool
+from vedana_core.settings import settings
+
+QueryResult = list[Record] | Exception
+
+
+# TODO replace with VTSArgs and CypherArgs
+class CypherQuery(str):
+    ...
+
+
+@dataclass
+class VTSQuery:
+    label: str
+    param: str
+    query: str
+
+    def __str__(self) -> str:
+        return f'vector_search("{self.label}","{self.param}","{self.query}")'
+
+
+DBQuery = CypherQuery | VTSQuery | str
+
+
+@dataclass
+class RagResults:
+    fts_res: list[Record] | None = None
+    vts_res: list[Record] | None = None
+    db_query_res: list[tuple[DBQuery, QueryResult]] | None = None
+
+
+class VTSArgs(BaseModel):
+    label: str = Field(description="node label")
+    property: str = Field(description="node property to search in")
+    text: str = Field(description="text to search similar")
+
+
+class CypherArgs(BaseModel):
+    query: str = Field(description="Cypher query")
+
+
+class GetHistoryArgs(BaseModel):
+    max_history: int = Field(20000, description="Maximum text length to retrieve from history. Cuts off on messages")
+
+
+VTS_TOOL_NAME = "vector_text_search"
+CYPHER_TOOL_NAME = "cypher"
+
+
+class RagAgent:
+    _vts_args: type[VTSArgs]
+
+    def __init__(
+        self,
+        graph: Graph,
+        vts: VectorStore,
+        data_model_description: str,
+        data_model_vts_indices,
+        llm: LLM,
+        ctx: ThreadContext,
+        logger: logging.Logger | None = None,
+    ) -> None:
+        self.graph = graph
+        self.vts = vts
+        self.llm = llm
+        self.logger = logger or logging.getLogger(__name__)
+        self._vts_meta_args: dict[str, dict[str, tuple[str, float]]] = {}  # stuff not passed through toolcall
+        self._vts_args = self._build_vts_arg_model(data_model_vts_indices)
+        self.data_model_description = data_model_description
+        self.ctx = ctx
+
+    def _build_vts_arg_model(self, vts_indices) -> Type[VTSArgs]:
+        """Create a Pydantic model with Enum-constrained fields for the VTS tool."""
+
+        if not vts_indices:
+            return VTSArgs
+
+        # fill in lookup for resolving idx type and getting custom threshold
+        for idx_type, i_name, i_attr, i_th in vts_indices:
+            if not self._vts_meta_args.get(i_name):
+                self._vts_meta_args[i_name] = {}
+            self._vts_meta_args[i_name][i_attr] = (idx_type, i_th)
+
+        # Label Enum – keys of `vts_indices`
+        LabelEnum = enum.Enum("LabelEnum", {name: name for (_type, name, _attr, _th) in vts_indices})  # type: ignore
+
+        # Property Enum – unique values of `vts_indices`
+        unique_props = set(attr for (_type, _name, attr, _th) in vts_indices)
+        prop_member_mapping: dict[str, str] = {}
+
+        used_names: set[str] = set()
+        for idx, prop in enumerate(sorted(unique_props)):
+            sanitized = re.sub(r"\W|^(?=\d)", "_", prop)
+            if sanitized in used_names:
+                sanitized = f"{sanitized}_{idx}"
+            used_names.add(sanitized)
+            prop_member_mapping[sanitized] = prop
+
+        PropertyEnum = enum.Enum("PropertyEnum", prop_member_mapping)  # type: ignore
+
+        VTSArgsEnum = create_model(
+            "VTSArgsEnum",
+            label=(LabelEnum, Field(description="node or edge label")),
+            property=(PropertyEnum, Field(description="property to search in")),
+            text=(str, Field(description="text for semantic search")),
+            __base__=VTSArgs,
+        )
+
+        return VTSArgsEnum
+
+    async def search_vector_text(
+        self,
+        label: str,
+        prop_type: str,
+        prop_name: str,
+        search_value: str,
+        threshold: float,
+        top_n: int = 5,
+    ) -> list[Record]:
+        embed = await self.llm.llm.create_embedding(search_value)
+        return await self.vts.vector_search(label, prop_type, prop_name, embed, threshold=threshold, top_n=top_n)
+
+    @staticmethod
+    def result_to_text(query: str, result: list[Record] | Exception) -> str:
+        if isinstance(result, Exception):
+            return f"Query: {query}\nResult: 'Error executing query'"
+        rows_str = "\n".join(row_to_text(row) for row in result)
+        return f"Query: {query}\nRows:\n{rows_str}"
+
+    async def execute_cypher_query(self, query, rows_limit: int = 30) -> QueryResult:
+        try:
+            return list(islice(await self.graph.execute_ro_cypher_query(query), rows_limit))
+        except Exception as e:
+            self.logger.exception(e)
+            return e
+
+    def rag_results_to_text(self, results: RagResults) -> str:
+        all_results = results.db_query_res or []
+        if results.vts_res:
+            all_results.append(("Vector text search", results.vts_res))
+        if results.fts_res:
+            all_results.append(("Full text search", results.fts_res))
+        return "\n\n".join(self.result_to_text(str(q), r) for q, r in all_results)
+
+    async def text_to_answer_with_vts_and_cypher(
+        self, text_query: str, threshold: float, top_n: int = 5
+    ) -> tuple[str, list, list[VTSQuery], list[CypherQuery]]:
+        vts_queries: list[VTSQuery] = []
+        cypher_queries: list[CypherQuery] = []
+
+        async def vts_fn(args: VTSArgs) -> str:
+            label = args.label.value if isinstance(args.label, enum.Enum) else args.label
+            prop = args.property.value if isinstance(args.property, enum.Enum) else args.property
+
+            # default fallback treats toolcall as node vector search
+            prop_type, th = self._vts_meta_args.get(label, {}).get(prop, ("node", threshold))
+            self.logger.debug(f"vts_fn(on={prop_type} label={label}, property={prop}, th={th}, n={top_n})")
+
+            vts_queries.append(VTSQuery(label, prop, args.text))
+            vts_res = await self.search_vector_text(label, prop_type, prop, args.text, threshold=th, top_n=top_n)
+            return self.result_to_text(VTS_TOOL_NAME, vts_res)
+
+        async def cypher_fn(args: CypherArgs) -> str:
+            self.logger.debug(f"cypher_fn({args})")
+            cypher_queries.append(CypherQuery(args.query))
+            res = await self.execute_cypher_query(args.query)
+            return self.result_to_text(CYPHER_TOOL_NAME, res)
+
+        vts_tool = Tool(
+            VTS_TOOL_NAME,
+            "Vector search for similar text in node properties, use for semantic search",
+            self._vts_args,
+            vts_fn,
+        )
+
+        cypher_tool = Tool(
+            CYPHER_TOOL_NAME,
+            "Execute a Cypher query against the graph database. Use for structured data retrieval and graph traversal.",
+            CypherArgs,
+            cypher_fn,
+        )
+
+        tools: list[Tool] = [vts_tool, cypher_tool]
+
+        all_query_events, answer = await self.llm.generate_cypher_query_with_tools(
+            data_descr=self.data_model_description,
+            messages=self.ctx.history[-settings.pipeline_history_length :],
+            tools=tools,
+        )
+
+        if not answer:
+            self.logger.warning(f"No answer found for {text_query}. Generating empty answer...")
+            answer = await self.llm.generate_no_answer(self.ctx.history[-settings.pipeline_history_length :])
+
+        return answer, all_query_events, vts_queries, cypher_queries
+
+
+def _remove_embeddings(val: Any):
+    if isinstance(val, Mapping):
+        return {k: _remove_embeddings(v) for k, v in val.items() if not k.endswith("_embedding")}
+    if isinstance(val, list):
+        return [_remove_embeddings(v) for v in val]
+    return val
+
+
+def _clear_record_val(val: Any):  # todo check for edges
+    params = _remove_embeddings(val)
+    if isinstance(val, neo4j.graph.Node) and isinstance(params, dict):
+        params["labels"] = list(val.labels)
+    return params
+
+
+def row_to_text(row: Any) -> str:
+    if isinstance(row, neo4j.Record):
+        row = {k: _clear_record_val(v) for k, v in row.items()}
+    try:
+        return json.dumps(row, ensure_ascii=False, indent=2)
+    except TypeError:
+        return str(row)

vedana_core/rag_pipeline.py

@@ -0,0 +1,326 @@
+import logging
+import traceback
+from dataclasses import asdict
+from typing import Any
+from pydantic import BaseModel, Field
+
+from jims_core.thread.thread_context import ThreadContext
+
+from vedana_core.data_model import DataModel
+from vedana_core.graph import Graph
+from vedana_core.vts import VectorStore
+from vedana_core.llm import LLM
+from vedana_core.rag_agent import RagAgent
+from vedana_core.settings import settings
+
+
+class DataModelSelection(BaseModel):
+    reasoning: str = Field(
+        default="",
+        description="Brief explanation of why these elements were selected for answering the user's question",
+    )
+    anchor_nouns: list[str] = Field(
+        default_factory=list,
+        description="List of anchor nouns (node types) needed to answer the question",
+    )
+    link_sentences: list[str] = Field(
+        default_factory=list,
+        description="List of link sentences (relationship types) needed to answer the question",
+    )
+    anchor_attribute_names: list[str] = Field(
+        default_factory=list,
+        description="List of anchor attribute names needed to answer the question (attributes belonging to nodes)",
+    )
+    link_attribute_names: list[str] = Field(
+        default_factory=list,
+        description="List of link attribute names needed to answer the question (attributes belonging to relationships)",
+    )
+    query_ids: list[str] = Field(
+        default_factory=list,
+        description="List of query scenario ID's that match the user's question type",
+    )
+
+
+dm_filter_base_system_prompt = """\
+Ты — помощник по анализу структуры графовой базы данных.
+
+Твоя задача: проанализировать вопрос пользователя и определить, какие элементы модели данных (узлы, связи, атрибуты, сценарии запросов) необходимы для формирования ответа.
+
+**Правила выбора:**
+1. Выбирай только те элементы, которые ДЕЙСТВИТЕЛЬНО нужны для ответа на вопрос
+2. Если вопрос касается связи между сущностями — выбери соответствующие узлы И связь между ними
+3. Выбирай атрибуты, которые могут содержать искомую информацию или использоваться для фильтрации
+   - anchor_attribute_names: атрибуты узлов (находятся в разделе "Атрибуты узлов")
+   - link_attribute_names: атрибуты связей (находятся в разделе "Атрибуты связей")
+4. Выбирай сценарий запроса, который лучше всего соответствует типу вопроса пользователя
+5. Лучше выбрать чуть больше, чем упустить важное — но не выбирай всё подряд
+
+**Формат ответа:**
+Верни JSON с выбранными элементами. Используй ТОЧНЫЕ имена из предоставленной модели данных.
+"""
+
+dm_filter_user_prompt_template = """\
+**Вопрос пользователя:**
+{user_query}
+
+**Модель данных:**
+{compact_data_model}
+
+Проанализируй вопрос и выбери необходимые элементы модели данных для формирования ответа.
+"""
+
+
+class StartPipeline:
+    """
+    Response for /start command
+    """
+
+    def __init__(self, data_model: DataModel) -> None:
+        self.data_model = data_model
+
+    async def __call__(self, ctx: ThreadContext) -> None:
+        lifecycle_events = await self.data_model.conversation_lifecycle_events()
+        start_response = lifecycle_events.get("/start")
+
+        message = start_response or "Bot online. No response for /start command in LifecycleEvents"
+        ctx.send_message(message)
+
+
+class RagPipeline:
+    """RAG Pipeline with data model filtering for optimized query processing.
+
+    This pipeline adds an initial step that filters the data model based on
+    the user's query, selecting only relevant anchors, links, attributes,
+    and query scenarios. This reduces token usage and improves LLM precision
+    for large data models.
+    """
+
+    def __init__(
+        self,
+        graph: Graph,
+        vts: VectorStore,
+        data_model: DataModel,
+        logger,
+        threshold: float = 0.8,
+        top_n: int = 5,
+        model: str | None = None,
+        filter_model: str | None = None,
+        enable_filtering: bool | None = None,
+    ):
+        self.graph = graph
+        self.vts = vts
+        self.data_model = data_model
+        self.logger = logger or logging.getLogger(__name__)
+        self.threshold = threshold
+        self.top_n = top_n
+        self.model = model or settings.model
+        self.filter_model = filter_model or settings.filter_model  # or self.model
+        self.enable_filtering = enable_filtering or settings.enable_dm_filtering
+
+    async def __call__(self, ctx: ThreadContext) -> None:
+        """Main pipeline execution - implements JIMS Pipeline protocol."""
+
+        # Get the last user message
+        user_query = ctx.get_last_user_message()
+        if not user_query:
+            ctx.send_message("I didn't receive a question. Please ask me something!")
+            return
+
+        try:
+            # Update status
+            await ctx.update_agent_status("Processing your question...")
+
+            # Process the query using RAG
+            answer, agent_query_events, technical_info = await self.process_rag_query(user_query, ctx)
+
+            # Send the answer
+            ctx.send_message(answer)
+
+            # Store technical information as an event
+            ctx.send_event(
+                "rag.query_processed",
+                {
+                    "query": user_query,
+                    "answer": answer,
+                    "technical_info": technical_info,
+                    "threshold": self.threshold,
+                },
+            )
+
+        except Exception as e:
+            self.logger.exception(f"Error in RAG pipeline: {e}")
+            error_msg = "An error occurred while processing the request"  # не передаем ошибку пользователю в диалог
+            ctx.send_message(error_msg)
+
+            # Store error event
+            ctx.send_event(
+                "rag.error",
+                {
+                    "query": user_query,
+                    "error": str(e),
+                    "traceback": traceback.format_exc(),
+                },
+            )
+
+    async def process_rag_query(self, query: str, ctx: ThreadContext) -> tuple[str, list, dict[str, Any]]:
+        # 1. Filter data model
+        if self.enable_filtering:
+            await ctx.update_agent_status("Analyzing query structure...")
+            data_model_description, filter_selection = await self.filter_data_model(query, ctx)
+        else:
+            data_model_description = await self.data_model.to_text_descr()
+            filter_selection = DataModelSelection()
+
+        # Read required DataModel properties
+        prompt_templates = await self.data_model.prompt_templates()
+        data_model_vector_search_indices = await self.data_model.vector_indices()
+
+        # 2. Create LLM and agent with filtered data model; Step 1 LLM costs are counted since same ctx.llm is used
+        llm = LLM(ctx.llm, prompt_templates=prompt_templates, logger=self.logger)
+        await ctx.update_agent_status("Searching knowledge base...")
+
+        if self.model != llm.llm.model and settings.debug:
+            llm.llm.set_model(self.model)
+
+        agent = RagAgent(
+            graph=self.graph,
+            vts=self.vts,
+            data_model_description=data_model_description,
+            data_model_vts_indices=data_model_vector_search_indices,
+            llm=llm,
+            ctx=ctx,
+            logger=self.logger,
+        )
+
+        (
+            answer,
+            agent_query_events,
+            vts_queries,
+            cypher_queries,
+        ) = await agent.text_to_answer_with_vts_and_cypher(
+            query,
+            threshold=self.threshold,
+            top_n=self.top_n,
+        )
+
+        technical_info: dict[str, Any] = {
+            "vts_queries": [str(q) for q in vts_queries],
+            "cypher_queries": [str(q) for q in cypher_queries],
+            "num_vts_queries": len(vts_queries),
+            "num_cypher_queries": len(cypher_queries),
+            "model_used": self.model,
+            "model_stats": {m: asdict(u) for m, u in ctx.llm.usage.items()},
+        }
+
+        # Add filtering info if applicable
+        if self.enable_filtering:
+            dm_anchors = await self.data_model.get_anchors()
+            dm_links = await self.data_model.get_links()
+            dm_queries = await self.data_model.get_queries()
+
+            # Count total attributes for original_counts
+            total_anchor_attrs = sum(len(anchor.attributes) for anchor in dm_anchors)
+            total_link_attrs = sum(len(link.attributes) for link in dm_links)
+
+            dm_filtering = {
+                "filter_model": self.filter_model,
+                "reasoning": filter_selection.reasoning,
+                "selected_anchors": filter_selection.anchor_nouns,
+                "selected_links": filter_selection.link_sentences,
+                "selected_anchor_attributes": filter_selection.anchor_attribute_names,
+                "selected_link_attributes": filter_selection.link_attribute_names,
+                "selected_queries": [dm_queries[int(i)].name for i in filter_selection.query_ids],
+                "original_counts": {
+                    "anchors": len(dm_anchors),
+                    "links": len(dm_links),
+                    "anchor_attrs": total_anchor_attrs,
+                    "link_attrs": total_link_attrs,
+                    "queries": len(dm_queries),
+                },
+                "filtered_counts": {
+                    "anchors": len(filter_selection.anchor_nouns),
+                    "links": len(filter_selection.link_sentences),
+                    "anchor_attrs": len(filter_selection.anchor_attribute_names),
+                    "link_attrs": len(filter_selection.link_attribute_names),
+                    "queries": len(filter_selection.query_ids),
+                },
+            }
+
+            ctx.send_event("rag.data_model_filtered", dm_filtering)
+
+        return answer, agent_query_events, technical_info
+
+    async def filter_data_model(
+        self,
+        query: str,
+        ctx: ThreadContext,
+    ) -> tuple[str, DataModelSelection]:
+        # Get description for filtering
+        dm_json = await self.data_model.to_compact_json()
+        dm_prompt_templates = await self.data_model.prompt_templates()
+
+        # Build the prompt
+        system_prompt = dm_prompt_templates.get("dm_filter_prompt", dm_filter_base_system_prompt)
+        user_prompt = dm_prompt_templates.get("dm_filter_user_prompt", dm_filter_user_prompt_template).format(
+            user_query=query,
+            compact_data_model=dm_json,
+        )
+
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ]
+
+        self.logger.debug(f"Filtering data model for query: {query}")
+
+        try:
+            filter_llm = ctx.llm
+
+            base_model = ctx.llm.model
+            if self.filter_model:  # if different model specified for filtering - use it
+                filter_llm.set_model(self.filter_model)
+
+            # Use structured output to get the selection
+            selection = await filter_llm.chat_completion_structured(messages, DataModelSelection)
+
+            if base_model:  # select base model back
+                ctx.llm.set_model(base_model)
+
+            if selection is None:
+                raise ValueError("LLM returned empty response")
+
+            # parse query id's to query names - LLM often misspells arbitrary query_names
+            query_names = [dm_json["queries"].get(int(i)) for i in selection.query_ids]
+
+            self.logger.debug(
+                f"Data model filter selection: "
+                f"anchors={selection.anchor_nouns}, "
+                f"links={selection.link_sentences}, "
+                f"anchor_attrs={selection.anchor_attribute_names}, "
+                f"link_attrs={selection.link_attribute_names}, "
+                f"queries={query_names}",
+            )
+            self.logger.debug(f"Filter reasoning: {selection.reasoning}")
+
+            # Create filtered data model
+            filtered_dm_descr = await self.data_model.to_text_descr(
+                anchor_nouns=selection.anchor_nouns,
+                link_sentences=selection.link_sentences,
+                anchor_attribute_names=selection.anchor_attribute_names,
+                link_attribute_names=selection.link_attribute_names,
+                query_names=query_names,
+            )
+
+            return filtered_dm_descr, selection
+
+        except Exception as e:  # return full data_model
+            self.logger.exception(f"Data model filtering failed: {e}. Using full data model.")
+            descr = await self.data_model.to_text_descr()
+            return descr, DataModelSelection(
+                reasoning=f"Filtering failed: {e}. Using full data model.",
+                anchor_nouns=[],
+                link_sentences=[],
+                anchor_attribute_names=[],
+                link_attribute_names=[],
+                query_ids=[],
+            )

vedana_core/settings.py

@@ -0,0 +1,35 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+class VedanaCoreSettings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="",
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+    grist_server_url: str  # default 'https://api.getgrist.com'
+    grist_api_key: str
+    grist_data_model_doc_id: str
+    grist_data_doc_id: str
+
+    debug: bool = False
+    model: str = "gpt-4.1"
+    enable_dm_filtering: bool = False  # pipeline selection (experimental)
+    filter_model: str = "gpt-5-mini"
+    judge_model: str = "gpt-5-mini"
+    embeddings_model: str = "text-embedding-3-large"
+    embeddings_dim: int = 1024
+
+    pipeline_history_length: int = 20
+
+    memgraph_uri: str
+    memgraph_user: str
+    memgraph_pwd: str
+
+
+settings = VedanaCoreSettings()  # type: ignore

vedana_core/start_pipeline.py

@@ -0,0 +1,17 @@
+from jims_core.thread.thread_context import ThreadContext
+from vedana_core.data_model import DataModel
+
+
+class StartPipeline:
+    """
+    Response for /start command
+    """
+
+    def __init__(self, data_model: DataModel) -> None:
+        self.data_model = data_model
+
+    async def __call__(self, ctx: ThreadContext) -> None:
+        lifecycle_events = await self.data_model.conversation_lifecycle_events()
+        start_response = lifecycle_events.get("/start")
+        message = start_response or "Bot online. No response for /start command in LifecycleEvents"
+        ctx.send_message(message)

vedana_core/utils.py

@@ -0,0 +1,31 @@
+def cast_dtype(value, value_name: str, dtype: str | None):
+    """
+    умнейшая конвертация строковых переменных по заданному в дата-модели типу данных
+    """
+
+    try:
+        if dtype == "int" or dtype == "float":
+            if isinstance(value, str):
+                if len(value.strip(",").split(",")) == 2:
+                    value = value.replace(",", ".")
+                elif len(value.strip(",").split(",")) > 2:  # assume "," is 000 delimeter
+                    value = value.replace(",", "")
+                elif len(value.strip(".").split(".")) > 2:  # assume "." is 000 delimeter
+                    value = value.replace(".", "")
+
+            value = float(value)
+            if dtype == "int":
+                value = int(value)
+
+        elif dtype == "bool":
+            return str(value).strip().lower() in ["1", "true", "да", "есть"]
+        elif dtype == "str":
+            return str(value)
+
+        else:
+            return value  # Fallback
+
+    except (ValueError, TypeError):
+        print(f"cast_dtype error for value {value_name} {value} (expected dtype {dtype})")
+
+    return value