HADR

{#---
description: A template for semantic text chunking tasks.
author: Pablo B., Marlon S.
---#}

{% set system %}
You are an expert linguist tool that semantically separates text.
Your task is to split the provided text into semantically meaningful
segments without altering the original content. Follow these
instructions carefully:

---Input---
- Preserve the original text exactly: wording, spelling, punctuation,
  whitespace, line breaks, and formatting must remain unchanged
  within each chunk.
- The text may include markdown, code blocks (including triple
  backticks), tables, LaTeX, visual-to-text transcription, or any
  other formats.
- If the input text contains the sequence <|>, replace every
  occurrence with [PIPE] before chunking to avoid confusion with
  the output delimiters.

---Task---
- Carefully break up the text into logical and semantically meaningful
  parts. Each chunk should represent a single, coherent idea or unit
  of meaning (such as a sentence, paragraph, or code block).
- Treat code blocks, tables, or other structured elements as atomic
  units where appropriate.
- Do not add any commentary, explanation, or preamble before or after
  the output.
- Do not add or remove any formatting characters such as triple
  backticks or newline characters inside the chunks.
- Return only the chunks concatenated together, each enclosed by the
  delimiters <|> with no spaces or newlines between chunks.

---Output Format---
<|>ChunkOne<|><|>ChunkTwo<|><|>LastChunk<|>
{% endset %}

{% set user %}
Semantically segment the following text into distinct meaning units,
where each chunk represents a single coherent idea or topic.

---Input Text---
{{ text }}
{% endset %}

{#---
description: A template for text summarization tasks.
author: Pablo B., Marlon S.
---#}

{% set system %}
You are an expert summary tool tasked with writing a concise and
comprehensive summary of the input text. Follow these instructions
carefully:

---Input---
- The text may include markdown, code blocks (including triple
  backticks and newline characters), tables, LaTeX, visual-to-text
  transcription, or any other formats. When summarizing, convert
  these elements into clear, descriptive writing rather than
  including raw formatting or code, unless the code or formatting
  is critical to understanding the content.

---Task---
- Summarize the input text by including all key details and
  important points without adding any information not present in
  the original text.
- Do not add any commentary, explanation, or preamble before or
  after the output.
- Keep the summary concise, ideally within 3-5 sentences.
{% endset %}

{% set user %}
Write a summary of the following text, including as many key
details as possible.

---Input Text---
{{ text }}
{% endset %}

CREATE TABLE IF NOT EXISTS artifacts (
    uuid    TEXT PRIMARY KEY,
    path    TEXT NOT NULL DEFAULT '',
    type    TEXT NOT NULL DEFAULT '',
    summary TEXT
);

CREATE TABLE IF NOT EXISTS branches (
    uuid          TEXT PRIMARY KEY,
    artifact_uuid TEXT NOT NULL REFERENCES artifacts(uuid),
    chunk         TEXT NOT NULL DEFAULT '',
    path          TEXT NOT NULL DEFAULT '',
    summary       TEXT
);

CREATE TABLE IF NOT EXISTS branch_embeddings (
    branch_uuid TEXT PRIMARY KEY REFERENCES branches(uuid),
    embedding   TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS leafs (
    uuid        TEXT PRIMARY KEY,
    branch_uuid TEXT NOT NULL REFERENCES branches(uuid),
    text        TEXT NOT NULL DEFAULT '',
    entities    TEXT
);

CREATE TABLE IF NOT EXISTS leaf_embeddings (
    leaf_uuid TEXT PRIMARY KEY REFERENCES leafs(uuid),
    embedding TEXT NOT NULL
);

from qdrant_client import QdrantClient
from tuutrag.utils import log_timestamp
from pathlib import Path

IMAGE = "qdrant/qdrant:latest"


class VectorDB():
    def __init__(self, port: int, host: str):
        self.client = self.connect(port, host)

    def connect(self, port: int, host: str):
        port = int(port)
        host = str(host)
        storage = str(Path(__file__).parent / "data/qdrant")
        try:
            client = QdrantClient(
                port=port, host=host, timeout=50, https=False
            )
            log_timestamp(f"Connected to http://{host}:{port}/dashboard")
            return client
        except Exception as e:
            raise RuntimeError(f"Error connecting to Qdrant: {e}")

    def create_collection(self, collection_name, vector_params):
        if self.client.collection_exists(
            collection_name=collection_name
        ):
            return
        self.client.create_collection(
            collection_name=collection_name,
            vectors_config=vector_params,
        )

    def upsert(self, collection_name: str, point):
        self.client.upsert(
            collection_name=collection_name,
            points=point,
            wait=True,
        )

{#---
description: A template for named entity extraction tasks.
author: Pablo B., Marlon S.
---#}

{% set system %}
You are an entity extraction tool. Extract entities only from the
section labeled "Input Text". Follow these instructions carefully:

---Input---
- The entire input text represents a meaningful idea (such as a
  sentence, paragraph, or code block).
- If the entity name is case-insensitive, capitalize the first
  letter of each significant word (title case). Ensure consistent
  naming across the entire extraction process.
- If the input text contains the sequence <|>, replace it
  with [PIPE] to avoid confusion with the entity delimiter.

---Task---
- Identify all meaningful entities only from the input text.
- Do not add any commentary, explanation, or preamble.
- Return only the identified entities in the following format,
  with each entity enclosed by <|> delimiters.

---Output Format---
<|>EntityOne<|><|>EntityTwo<|><|>EntityThree<|>

---Example---
Input: "On March 14, 2023, Dr. Evelyn Hart of the Johns Hopkins
Neurology Department filed Case No. 2023-MED-0047 with the FDA
citing adverse reactions to Compound XR-9B, developed under NIH
Grant #R01NS112233, affecting patients in the Phase III TriNeuro
Trial across facilities in Baltimore, MD and Denver, CO."
Output: <|>March 14, 2023<|><|>Dr. Evelyn Hart<|>
<|>Johns Hopkins Neurology Department<|><|>Johns Hopkins<|>
<|>Case No. 2023-MED-0047<|><|>FDA<|><|>Compound XR-9B<|>
<|>NIH Grant #R01NS112233<|><|>NIH<|>
<|>Phase III TriNeuro Trial<|><|>Baltimore<|><|>MD<|>
<|>Denver<|><|>CO<|>
{% endset %}

{% set user %}
Extract all named entities from the following text.

---Input Text---
{{ text }}
{% endset %}

{#---
description: A template for local relationship identification
             tasks between entities.
author: Pablo B., Marlon S.
---#}

{% set system %}
You are a relationship identification tool. Identify relationships
only from the section labeled "Entities". "Raw Chunk" may be used
for disambiguation only. Follow these instructions carefully:

---Input---
- Entities are provided as a list, where each list item represents
  a single coherent idea.
- If the input text contains the sequence <|>, replace it
  with [PIPE] to avoid confusion with the entity delimiter.

---Task---
- Identify meaningful relationships between entities based on
  the provided information. Only extract relationships that are
  explicitly stated or directly implied.
- If a single statement describes a relationship involving more
  than two entities (an N-ary relationship), decompose it into
  multiple binary (two-entity) pairs.
  E.g. "REQ-12 and REQ-13 both derive from NEED-04" yields:
  "REQ-12, derives-from, NEED-04" and
  "REQ-13, derives-from, NEED-04."
- Relationships are directional: source → target.
- Do not add any commentary or preamble.

---Output Format---
<|>Source, relationship-type, Target<|>
{% endset %}

{% set user %}
Identify all relationships that are explicitly stated or directly
implied across the "Entities".

---Raw Chunk---
{{ raw_chunk }}

---Entities---
{{ entities }}
{% endset %}

{#---
description: A template for global relationship identification
             tasks between entities across multiple sets.
             This template is identical to relation_global.j2.
             Both map to the same underlying prompt.
author: Pablo B., Marlon S.
---#}

{% set system %}
You are a relationship identification tool. Identify relationships
only from the sections labeled "Entities". Each "Relations",
"Raw Chunk" and "Artifact Summary" may be used for disambiguation
of their respective set only. Follow these instructions carefully:

---Input---
Input is provided across K numbered sets, each formatted as:
 [SET {n}]
 Entities: <comma-separated list>
 Relations: <"<|>"-comma separated list>
 Raw Chunk: <source text>
 Artifact Summary: <artifact-level summary>
- Treat each entity as scoped to its own set.
- If the input text contains the sequence <|>, replace it
  with [PIPE].

---Task---
- Identify meaningful relationships between entities across
  different sets. Do not extract relationships within the same set.
- Only extract relationships explicitly stated or directly implied.
- Decompose N-ary relationships into binary pairs.
- Relationships are directional: source → target.

---Output Format---
<|>Source, relationship-type, Target<|>

---Example---
Input:
[SET 1]
Entities: Dr. Evelyn Hart, Johns Hopkins, Case No. 2023-MED-0047, FDA
Relations: <|>Dr. Evelyn Hart, works at, Johns Hopkins<|>
           <|>Case No. 2023-MED-0047, reviewed by, FDA<|>
Raw Chunk: "…filed Case No. 2023-MED-0047 with the FDA citing
           adverse reactions to Compound XR-9B."
Artifact Summary: "FDA case filing regarding drug adverse events."
[SET 2]
Entities: Compound XR-9B, NIH, Phase III TriNeuro Trial
Relations: <|>Compound XR-9B, funded by, NIH Grant #R01NS112233<|>
           <|>Compound XR-9B, tested in, Phase III TriNeuro Trial<|>
Raw Chunk: "Compound XR-9B…was the subject of an adverse event report…
           The Phase III TriNeuro Trial was subsequently suspended."
Artifact Summary: "NIH-funded clinical trial documentation."
Output:
<|>Case No. 2023-MED-0047, concerns, Compound XR-9B<|>
<|>Dr. Evelyn Hart, reported, Compound XR-9B<|>
<|>FDA, suspended, Phase III TriNeuro Trial<|>
{% endset %}

{% set user %}
Identify all relationships that are explicitly stated or directly
implied across the "Entities".

---Raw Chunk---
{{ raw_chunk }}

---Entities---
{{ entities }}

---Relations---
{{ relations }}

---Artifact Summary---
{{ artifact_summary }}
{% endset %}

import json
from typing import List
from exports import export
from neo4j import GraphDatabase


@export
class MemgraphConnection:
    obj_name = "MemgraphConnection"
    workspace = "kg"

    def __init__(self, port: int, frontend_port: int, host: str):
        self.driver = self.__connect(port, frontend_port, host)

    def __connect(self, port, frontend_port, host):
        URI = f"bolt://{host}:{int(port)}"
        driver = GraphDatabase.driver(URI, auth=("", ""))
        driver.verify_connectivity()
        print(f"Memgraph UI → http://{host}:{frontend_port}/")
        print(f"Memgraph Bolt → {URI}")
        return driver

    def read_data(self, file_path: str) -> List:
        parts = []
        with open(file_path, "r", encoding="utf-8") as f:
            for line in f:
                data = json.loads(line)
                content = data["text"].split("<|>")[1]
                part = [i.strip() for i in content.split(",")]
                if len(part) == 3:
                    parts.append(part)
        return parts

    def upsert(self, data: List) -> None:
        with self.driver.session() as session:
            for i, part in enumerate(data):
                session.run("""
                    MERGE (source:Entity {name: $source})
                    MERGE (target:Entity {name: $target})
                    MERGE (source)-[r:RELATIONSHIP]->(target)
                    SET r.type = $relationship
                """, source=part[0],
                     relationship=part[1],
                     target=part[2])
                print(f"upserted {i+1}/{len(data)}: "
                      f"{part[0]} -[{part[1]}]-> {part[2]}")

from typing import Literal, TypedDict


class BranchChunk(TypedDict):
    uuid: str
    chunk: str
    path: str
    type: str


class TextContent(TypedDict):
    type: Literal["text"]
    text: str


class Message(TypedDict):
    role: Literal["system", "user"]
    content: str | list[TextContent]


class RequestBody(TypedDict):
    model: str
    messages: list[Message]
    stream: Literal[False]


class BatchRequest(TypedDict):
    custom_id: str
    method: Literal["POST"]
    url: Literal["/v1/chat/completions"]
    body: RequestBody

from tuutrag.types import BatchRequest


def create_batch_request(
    custom_id: str,
    model: str,
    **kwargs
) -> BatchRequest:
    return {
        "custom_id": custom_id,
        "method": "POST",
        "url": "/v1/chat/completions",
        "body": {
            "model": model,
            "messages": [
                {"role": "system", "content": kwargs["system"]},
                {"role": "user",   "content": kwargs["user"]},
            ],
            "stream": False,
        },
    }

from pathlib import Path
from jinja2 import Environment, FileSystemLoader, meta


TEMPLATES_DIR = Path(__file__).parent / "templates"


def load_prompt(template_name: str, **kwargs) -> dict[str, str]:
    env = Environment(
        loader=FileSystemLoader(str(TEMPLATES_DIR)),
        keep_trailing_newline=True,
        comment_start_string="{#---",
        comment_end_string="---#}",
    )
    template = env.get_template(template_name)
    module = template.make_module(vars=kwargs)
    return {
        "system": str(module.system).strip(),
        "user":   str(module.user).strip(),
    }

from tiktoken.core import Encoding
from tuutrag.types import BatchRequest


def count_batch_request_tokens(
    enc: Encoding, payload: BatchRequest
) -> int:
    total = 0
    for message in payload["body"]["messages"]:
        content = message["content"]
        if isinstance(content, str):
            total += len(enc.encode(content))
        elif isinstance(content, list):
            for block in content:
                if block.get("type") == "text":
                    total += len(enc.encode(block["text"]))
    return total