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

vedana_core/graph.py

@@ -0,0 +1,300 @@
+import abc
+import json
+import logging
+import re
+from typing import Any, Dict, Iterable, Set, cast
+
+import aioitertools as aioit
+import neo4j
+import numpy as np
+import typing_extensions as te
+from neo4j import AsyncGraphDatabase, EagerResult, RoutingControl
+from opentelemetry import trace
+
+logger = logging.getLogger(__name__)
+tracer = trace.get_tracer(__name__)
+
+Record = neo4j.Record
+
+
+class Graph(abc.ABC):
+    async def add_node(
+        self,
+        node_id: str,
+        labels: Set[str],
+        properties: dict[str, Any] | None = None,
+        embeddings: dict[str, np.ndarray] | None = None,
+    ) -> None:
+        raise NotImplementedError
+
+    async def add_edge(self, from_id: str, to_id: str, type_: str, attrs: Dict[str, Any] | None) -> None:
+        raise NotImplementedError
+
+    async def number_of_nodes(self) -> int:
+        raise NotImplementedError
+
+    async def number_of_edges(self) -> int:
+        raise NotImplementedError
+
+    async def run_cypher(
+        self,
+        query: str,
+        parameters: dict[str, Any] | None = None,
+        limit: int | None = None,
+    ) -> Iterable[Record]:
+        raise NotImplementedError
+
+    async def get_existing_node_types(self) -> Iterable[list[str]]:
+        raise NotImplementedError
+
+    async def llm_schema(self) -> str:
+        raise NotImplementedError
+
+    async def text_search(self, label: str, query: str, limit: int = 10) -> Iterable[Record]:
+        raise NotImplementedError
+
+    async def setup(self, *_, create_basic_indices: bool = True, **kwargs) -> None:
+        # Set false to speedup import
+        if create_basic_indices:
+            await self.create_basic_indices()
+
+    async def create_basic_indices(self) -> None:
+        ...
+
+    async def execute_ro_cypher_query(
+        self,
+        query: str,
+        parameters: dict[str, Any] | None = None,
+        limit: int | None = None,
+    ) -> Iterable[Record]:
+        return await self.run_cypher(query, parameters, limit=limit)
+
+    async def clear(self) -> None:
+        ...
+
+    def close(self) -> None:
+        ...
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+class CypherGraph(Graph):
+    async def add_node(
+        self,
+        node_id: str,
+        labels: Set[str],
+        properties: dict[str, Any] | None = None,
+        embeddings: dict[str, np.ndarray] | None = None,
+    ) -> None:
+        query, params = self._add_node_cypher(node_id, labels, properties or {})
+        await self.run_cypher(query, params)
+
+    def _add_node_cypher(
+        self,
+        node_id: str,
+        labels: Set[str],
+        properties: dict[str, Any],
+    ) -> tuple[str, dict[str, Any]]:
+        labels_expr = escape_labels(labels)
+        props = {
+            **properties,
+            "id": node_id,
+        }
+        # TODO escape
+        pros_expr = ", ".join(f"{k}: ${k}" for k in props.keys())
+        return (
+            f"MERGE (n:{labels_expr} {{id: $id}}) SET n = {{{pros_expr}}} RETURN n",
+            props,
+        )
+
+    async def add_edge(self, from_id: str, to_id: str, type_: str, attrs: Dict[str, Any] | None) -> None:
+        query, params = self._add_edge_cypher(from_id, to_id, type_, attrs)
+        await self.run_cypher(query, params)
+
+    def _add_edge_cypher(
+        self, from_id: str, to_id: str, type_: str, attrs: Dict[str, Any] | None
+    ) -> tuple[str, dict[str, Any]]:
+        attrs = attrs or {}
+        labels_expr = escape_labels({type_})
+        # attrs = {escape_cypher(k): v for k, v in attrs.items()}
+        attrs_expr = ", ".join(f"{k}: ${k}" for k in attrs.keys() if k)
+        params = {
+            **attrs,
+            "from_id": from_id,
+            "to_id": to_id,
+        }
+        return (
+            "MATCH (nf {id: $from_id}), (nt {id: $to_id}) "
+            f"CREATE (nf)-[r:{labels_expr} {{{attrs_expr}}}]->(nt) RETURN r",
+            params,
+        )
+
+    async def add_edges(self, edges: Iterable[tuple[str, str, dict]], **common_attrs) -> None:
+        for edge_tuple in edges:
+            from_id, to_id, attrs = edge_tuple
+            attrs = {**common_attrs, **attrs}
+            labels: Iterable[str] = attrs.pop("__labels__", [])
+            type_ = next(iter(labels), "no_type")
+            await self.add_edge(from_id, to_id, type_, attrs)
+
+    async def number_of_nodes(self) -> int:
+        res = await self.execute_ro_cypher_query("MATCH (n) RETURN count(*) as cnt")
+        return next(iter(res))["cnt"]
+
+    async def number_of_edges(self) -> int:
+        res = await self.execute_ro_cypher_query("MATCH (f)-[]->(t) RETURN count(*) as cnt")
+        return next(iter(res))["cnt"]
+
+    async def get_existing_node_types(self) -> Iterable[list[str]]:
+        res = await self.execute_ro_cypher_query("MATCH (n) RETURN DISTINCT labels(n) as l;")
+        return [r["l"] for r in res]
+
+
+# class NXGraph(Graph):
+#     def __init__(self, graph: nx.Graph) -> None:
+#         self.graph: nx.Graph = graph
+#         self.gcypher = GrandCypher(self.graph)
+
+#     def execute_ro_cypher_query(self, query: str) -> Iterable[Any]:
+#         return self.gcypher.run(query)
+
+#     def add_node(self, node_id: str, labels: Set[str], **attributes) -> None:
+#         self.graph.add_node(node_id, __labels__=labels, **attributes)
+
+#     def number_of_edges(self) -> int:
+#         return self.graph.number_of_edges()
+
+#     def clear(self) -> None:
+#         self.graph.clear()
+
+
+class MemgraphGraph(CypherGraph):
+    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 execute_ro_cypher_query(
+        self, query: str, parameters: dict[str, Any] | None = None, limit: int | None = None
+    ) -> Iterable[Record]:
+        with tracer.start_as_current_span("memgraph.execute_ro_cypher_query") as span:
+            span.set_attribute("memgraph.query", query)
+            if parameters:
+                span.set_attribute("memgraph.parameters", json.dumps(parameters))
+            result: EagerResult = await self.driver.execute_query(query, parameters, routing_=RoutingControl.READ)
+
+        return result.records
+
+    async def run_cypher(
+        self,
+        query: str,
+        parameters: dict[str, Any] | None = None,
+        limit: int | None = None,
+    ) -> Iterable[Record]:
+        with tracer.start_as_current_span("memgraph.run_cypher") as span:
+            span.set_attribute("memgraph.query", query)
+            if parameters:
+                span.set_attribute("memgraph.parameters", json.dumps(parameters))
+            if limit is not None:
+                span.set_attribute("memgraph.limit", limit)
+
+            async with self.driver.session() as session:
+                result = await aioit.list(aioit.islice(await session.run(query, parameters), limit))
+
+        return result
+
+    async def add_node(
+        self,
+        node_id: str,
+        labels: Set[str],
+        properties: Dict[str, Any] | None = None,
+        embeddings: Dict[str, np.ndarray] | None = None,
+    ) -> None:
+        if properties and embeddings:
+            embed_props = {f"{prop_name}_embedding": v for prop_name, v in embeddings.items()}
+            properties = {
+                **properties,
+                **embed_props,
+            }
+        await super().add_node(node_id, labels, properties, embeddings)
+
+    async def llm_schema(self) -> str:
+        """can be used as fallback data model structure"""
+        res = await self.driver.execute_query("CALL llm_util.schema() YIELD schema RETURN schema")
+        return res.records[0]["schema"]
+
+    async def create_basic_indices(self, node_types=None) -> None:
+        if not node_types:
+            node_types = await self.get_existing_node_types()
+        for label in node_types:
+            await self.create_node_prop_index(set(label), "id", unique=True)
+
+    async def clear(self) -> None:
+        async with self.driver.session() as session:
+            res = await session.run("CALL vector_search.show_index_info() YIELD index_name RETURN *")
+
+            async for (idx_name,) in res:
+                await session.run(f"DROP VECTOR INDEX {escape_cypher(idx_name)}")
+            idx_name_re = re.compile(r"\(name:\s(.+?)\)")
+            async for row in await session.run(cast(te.LiteralString, "SHOW INDEX INFO")):
+                index_type = row["index type"]
+                idx_name = next(iter(idx_name_re.findall(index_type)), None)
+                if not idx_name:
+                    continue
+                await session.run(f"DROP TEXT INDEX {escape_cypher(idx_name)}")
+            await session.run("CALL schema.assert({}, {}, {}, true) YIELD action, key, keys, label, unique")
+            await session.run("MATCH (n) DETACH DELETE n")
+            # TODO more efficient:
+            # USING PERIODIC COMMIT num_rows
+            # MATCH (n)-[r]->(m)
+            # DELETE r;
+            # USING PERIODIC COMMIT num_rows
+            # MATCH (n)
+            # DETACH DELETE n;
+
+    async def text_search(self, label: str, query: str, limit: int = 10) -> Iterable[Record]:
+        with tracer.start_as_current_span("memgraph.text_search") as span:
+            span.set_attribute("memgraph.label", label)
+            span.set_attribute("memgraph.fts_query", query)
+            span.set_attribute("memgraph.limit", limit)
+
+            query = "CALL text_search.search_all($idx_name, $query) YIELD node RETURN node LIMIT $limit"
+            span.set_attribute("memgraph.query", query)
+
+            res = await self.driver.execute_query(
+                query,
+                idx_name=self._fts_idx_name(label),
+                query=query,
+                limit=limit,
+                routing_=RoutingControl.READ,
+            )
+            return res.records
+
+    async def create_node_prop_index(self, labels: set[str], property: str, unique: bool = False) -> None:
+        escaped_label = escape_labels(labels)
+        escaped_prop = escape_cypher(property)
+        await self.run_cypher(f"CREATE INDEX ON :{escaped_label}({escaped_prop})")
+        if not unique:
+            return
+        await self.run_cypher(f"CREATE CONSTRAINT ON (n:{escaped_label})\nASSERT n.{escaped_prop} IS UNIQUE")
+
+    @staticmethod
+    def _fts_idx_name(label: str) -> str:
+        return f"{label.lower()}_fts_idx"
+
+    def close(self):
+        self.driver.close()
+
+
+def escape_cypher(identifier: str) -> str:
+    identifier = identifier.replace("\u0060", "`").replace("`", "``")
+    return f"`{identifier}`"
+
+
+def escape_labels(labels: set[str]) -> str:
+    return ":".join(escape_cypher(label) for label in labels)

vedana_core/llm.py

@@ -0,0 +1,192 @@
+import asyncio
+import logging
+from typing import Awaitable, Callable, Iterable
+
+import openai
+from jims_core.llms.llm_provider import LLMProvider
+from jims_core.thread.schema import CommunicationEvent
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionToolMessageParam,
+)
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class Tool[T: BaseModel]:
+    def __init__(
+        self, name: str, description: str, args_cls: type[T], fn: Callable[[T], Awaitable[str]] | Callable[[T], str]
+    ) -> None:
+        self.name = name
+        self.description = description
+        self.args_cls = args_cls
+        self.fn = fn
+        self.openai_def = openai.pydantic_function_tool(args_cls, name=name, description=description)
+
+    async def call(self, args_json: str) -> str:
+        try:
+            fn_args = self.args_cls.model_validate_json(args_json)
+        except ValueError:
+            return f"Invalid tool args: {args_json}"
+
+        if asyncio.iscoroutinefunction(self.fn):
+            result = await self.fn(fn_args)
+        else:
+            result: str = await asyncio.to_thread(self.fn, fn_args)  # type: ignore
+
+        return result
+
+
+class LLM:
+    def __init__(
+        self,
+        llm_provider: LLMProvider,
+        prompt_templates: dict[str, str],
+        logger: logging.Logger | None = None,
+    ) -> None:
+        self.logger = logger or logging.getLogger(__name__)
+        self.llm = llm_provider
+        self.prompt_templates = prompt_templates
+
+    # Current
+    async def generate_cypher_query_with_tools(
+        self,
+        data_descr: str,
+        messages: Iterable,
+        tools: list[Tool],
+    ) -> tuple[list[ChatCompletionMessageParam], str]:
+        tool_names = [t.name for t in tools]
+        msgs = make_cypher_query_with_tools_dialog(data_descr, self.prompt_templates, messages, tool_names=tool_names)
+        return await self.create_completion_with_tools(msgs, tools=tools)
+
+    async def create_completion_with_tools(
+        self,
+        messages: list[ChatCompletionMessageParam],
+        tools: Iterable[Tool],
+    ) -> tuple[list[ChatCompletionMessageParam], str]:
+        messages = messages.copy()
+        tool_defs = [tool.openai_def for tool in tools]
+        tools_map = {tool.name: tool for tool in tools}
+
+        async def _execute_tool_call(tool_call):
+            tool_name = tool_call.function.name
+            tool = tools_map.get(tool_name)
+            if not tool:
+                self.logger.error(f"Tool {tool_name} not found!")
+                return tool_call.id, f"Tool {tool_name} not found!"
+
+            self.logger.debug(f"Calling tool {tool_name}")
+            try:
+                tool_res = await tool.call(tool_call.function.arguments)
+            except Exception as e:
+                self.logger.exception("Error executing tool %s: %s", tool_name, e)
+                tool_res = f"Error executing tool {tool_name}: {e}"
+
+            self.logger.debug("Tool %s (%s) result: %s", tool_name, tool.description, tool_res)
+            return tool_call.id, tool_res
+
+        max_iters = 5
+        for i in range(max_iters):
+            msg, tool_calls = await self.llm.chat_completion_with_tools(
+                messages=messages,
+                tools=tool_defs,
+            )
+
+            messages.append(msg.to_dict())  # type: ignore
+
+            if not tool_calls:
+                self.logger.debug("No tool calls found. Exiting tool call loop")
+                break
+
+            self.logger.debug(f"Tool call iter {i + 1}/{max_iters}")
+
+            # Execute tool calls in parallel
+            results = await asyncio.gather(*[_execute_tool_call(t) for t in tool_calls])
+
+            for tool_call_id, tool_res in results:
+                messages.append(
+                    ChatCompletionToolMessageParam(role="tool", tool_call_id=tool_call_id, content=tool_res)
+                )
+
+            if i == max_iters - 1:
+                self.logger.warning(f"Reached tool call iteration limit ({max_iters}). Exiting tool call loop")
+                finalize_prompt = self.prompt_templates.get("finalize_answer_tmplt", finalize_answer_tmplt)
+                finalize_msg = {"role": "system", "content": finalize_prompt}
+                final_msg = await self.llm.chat_completion_plain(messages + [finalize_msg])
+                messages.append(final_msg.to_dict())  # type: ignore
+                break
+
+        for last_msg in reversed(messages):  # sometimes message with final answer is not the last one
+            if last_msg.get("role", "") == "assistant" and last_msg.get("content"):
+                return messages, str(last_msg.get("content"))
+        return messages, ""
+
+    async def generate_no_answer(
+        self,
+        dialog: list[CommunicationEvent] | None = None,
+    ) -> str:
+        prompt = self.prompt_templates.get("generate_no_answer_tmplt", generate_no_answer_tmplt)
+        messages = [
+            {"role": "system", "content": prompt},
+            *(dialog or []),
+        ]
+        response = await self.llm.chat_completion_plain(messages)
+        human_answer = "" if response.content is None else response.content.strip()
+        self.logger.debug(f"Generated 'no answer' response: {human_answer}")
+        return human_answer
+
+
+finalize_answer_tmplt = """\
+Сформулируй ответ на запрос пользователя основе информации, полученной в результате вызова результатов инструментов.
+Если информации недостаточно для точного ответа, ясно опиши ограничения и предложи 1–2 уточняющих вопроса.
+Важно! Не упоминай инструменты в явном виде, ссылайся только на данные.
+"""
+
+generate_no_answer_tmplt = """\
+Ты - помощник, который преобразует технические ответы в понятный человеку текст.
+Мы не смогли найти ответ на вопрос пользователя в базе знаний.
+Сформулируй ответ, сообщающий кратко и информативно, что ответа не найдено.
+Предложи пару вариантов уточняющих вопросов на основе информации в контексте. Предложи в casual стиле.
+"""
+
+generate_answer_with_tools_tmplt = """\
+Ты — помощник по работе с графовыми базами данных, в которых используется язык запросов Cypher
+
+Цель: постараться найти ответ на вопрос пользователя используя инструменты для работы с БД на основе текстового описания графовой базы данных.
+
+На вход ты получаешь graph_composition: – описание графа и примеры запросов по нему, и user_query – пользовательский запрос.
+
+**Что нужно сделать:**
+1. Сгенерировать `Cypher`-запросы, используя узлы, атрибуты и связи перечисленные в **graph_composition**.
+2. Руководствуйся данными в **graph_composition** примерами запросов, чтобы составить итоговый запрос.
+3. Используй инструменты {tools} для выполнения запросов и поиска
+
+Если нужно, используй несколько `MATCH`-блоков, например:
+    MATCH (o:offer)-[:OFFER_belongs_to_CATEGORY]->(c:category)
+    MATCH (o)-[:OFFER_made_of_MATERIAL]->(m:material)
+    WHERE c.category_name = "Встраиваемый светильник" AND m.material_name IN ["Стекло", "Металл и Стекло", "Алюминий и стекло"]
+    RETURN o
+
+Теперь проанализируй следующую структуру графа, и постарайся найти ответ на вопрос используя инструменты {tools}. (Лучше использовать несколько инструментов)
+
+**graph_composition**
+{graph_description}
+"""
+
+
+def make_cypher_query_with_tools_dialog(
+    graph_description: str,
+    prompt_templates: dict[str, str],
+    messages: Iterable,
+    tool_names: list[str],
+) -> list[ChatCompletionMessageParam]:
+    prompt_template = prompt_templates.get("generate_answer_with_tools_tmplt", generate_answer_with_tools_tmplt)
+    prompt = prompt_template.format(graph_description=graph_description, tools=", ".join(tool_names))
+    return [
+        {
+            "role": "system",
+            "content": prompt,
+        },
+        *messages,
+    ]