protoprompt 0.1.0__py3-none-any.whl

protoprompt/__init__.py

@@ -0,0 +1,58 @@
+"""protoprompt: layered context builder for LLM prompts.
+
+Public top-level exports live here; per-layer subpackages re-export
+their own public API as well, so both styles work:
+
+    from protoprompt import ContextBuilder
+    from protoprompt.context import ContextInput
+"""
+
+from protoprompt.context import ContextInput, ContextOutput
+from protoprompt.exceptions import TokenBudgetExceededError
+from protoprompt.injector import ContextBuilder
+from protoprompt.injector_budgeted import (
+    DEFAULT_PRIORITIES,
+    BudgetReport,
+    TokenBudgetedContextBuilder,
+)
+from protoprompt.llm import LLMClientProtocol
+from protoprompt.pipeline import Pipeline
+from protoprompt.profile.builder import ProfileBuilder
+from protoprompt.profile.types import UserProfile
+from protoprompt.session.compressor import Compressor
+from protoprompt.session.strategy import (
+    HeuristicStrategy,
+    LLMSummaryStrategy,
+    StrategyProtocol,
+)
+from protoprompt.session.types import CompressedBlock, Session
+from protoprompt.store.memory import InMemStore
+from protoprompt.store.protocol import StoreProtocol
+from protoprompt.tokens.protocol import TokenCounter
+from protoprompt.tokens.regex_counter import RegexTokenCounter
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "ContextBuilder",
+    "TokenBudgetedContextBuilder",
+    "ContextInput",
+    "ContextOutput",
+    "BudgetReport",
+    "DEFAULT_PRIORITIES",
+    "TokenBudgetExceededError",
+    "Pipeline",
+    "Compressor",
+    "StrategyProtocol",
+    "HeuristicStrategy",
+    "LLMSummaryStrategy",
+    "Session",
+    "CompressedBlock",
+    "UserProfile",
+    "ProfileBuilder",
+    "StoreProtocol",
+    "InMemStore",
+    "TokenCounter",
+    "RegexTokenCounter",
+    "LLMClientProtocol",
+]

protoprompt/context.py

@@ -0,0 +1,33 @@
+"""Public dataclasses for context assembly."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from protoprompt.injector_budgeted import BudgetReport
+
+
+@dataclass
+class ContextInput:
+    query: str
+    chat_id: str = ""
+    system_prompt: str = ""
+    doc_ids: list[int] = field(default_factory=list)
+    embedding_model: str = "nomic-embed-text"
+    top_k_rag: int = 5
+    top_k_session: int = 3
+    include_rag: bool = True
+    include_session: bool = True
+    include_profile: bool = False
+    profile_text: str = ""
+
+
+@dataclass
+class ContextOutput:
+    system_prompt: str
+    rag_blocks: list[str] = field(default_factory=list)
+    session_blocks: list[str] = field(default_factory=list)
+    profile_used: bool = False
+    budget_report: "BudgetReport | None" = None

protoprompt/exceptions.py

@@ -0,0 +1,20 @@
+"""Exceptions raised by the context assembly pipeline."""
+
+from __future__ import annotations
+
+
+class TokenBudgetExceededError(RuntimeError):
+    """Raised when a hard-required section (system prompt) does not fit
+    into the configured token budget.
+
+    Soft sections (RAG, session) are dropped silently by the budget
+    allocator; only mandatory sections ever trigger this.
+    """
+
+    def __init__(self, used: int, budget: int, section: str) -> None:
+        super().__init__(
+            f"Section '{section}' needs {used} tokens but budget is {budget}"
+        )
+        self.used = used
+        self.budget = budget
+        self.section = section

protoprompt/injector.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import logging
+
+from protoprompt.context import ContextInput, ContextOutput
+from protoprompt.llm import LLMClientProtocol
+from protoprompt.store.protocol import StoreProtocol
+
+logger = logging.getLogger(__name__)
+
+
+class ContextBuilder:
+    """Default context assembler.
+
+    RAG blocks (if any) are queried, session memory is queried when a
+    ``chat_id`` is supplied, and the system prompt is the anchor. The
+    query is embedded once and reused for both retrievals.
+    """
+
+    def __init__(
+        self,
+        store: StoreProtocol,
+        llm: LLMClientProtocol,
+    ) -> None:
+        self._store = store
+        self._llm = llm
+
+    async def build(self, inp: ContextInput) -> ContextOutput:
+        parts: list[str] = [inp.system_prompt] if inp.system_prompt else []
+        rag_blocks: list[str] = []
+        session_blocks: list[str] = []
+        profile_used = False
+
+        query_emb: list[float] | None = None
+        if (inp.include_rag and inp.doc_ids) or (inp.include_session and inp.chat_id):
+            query_emb = (await self._llm.embed([inp.query], model=inp.embedding_model))[0]
+
+        if inp.include_rag and inp.doc_ids and query_emb is not None:
+            str_doc_ids = [str(d) for d in inp.doc_ids]
+            where = {"doc_id": {"$in": str_doc_ids}} if len(str_doc_ids) > 1 else {"doc_id": str_doc_ids[0]}
+            results = self._store.query(query_emb, top_k=inp.top_k_rag, where=where)
+            if results:
+                rag_texts = [r["document"] for r in results]
+                rag_blocks = rag_texts
+                parts.append("\n\n---\n\n".join(rag_texts))
+
+        if inp.include_session and inp.chat_id and query_emb is not None:
+            session_results = self._store.query(
+                query_emb,
+                top_k=inp.top_k_session,
+                where={"doc_id": f"session_{inp.chat_id}"},
+            )
+            if session_results:
+                session_texts = [r["document"] for r in session_results]
+                session_blocks = session_texts
+                parts.append(
+                    "История диалога (сжатая):\n" + "\n---\n".join(session_texts)
+                )
+
+        if inp.include_profile and inp.profile_text:
+            parts.append(f"Профиль пользователя:\n{inp.profile_text}")
+            profile_used = True
+
+        return ContextOutput(
+            system_prompt="\n\n".join(parts),
+            rag_blocks=rag_blocks,
+            session_blocks=session_blocks,
+            profile_used=profile_used,
+        )

protoprompt/injector_budgeted.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+
+from protoprompt.context import ContextInput, ContextOutput
+from protoprompt.exceptions import TokenBudgetExceededError
+from protoprompt.injector import ContextBuilder
+from protoprompt.llm import LLMClientProtocol
+from protoprompt.store.protocol import StoreProtocol
+from protoprompt.tokens.protocol import TokenCounter
+from protoprompt.tokens.regex_counter import RegexTokenCounter
+
+logger = logging.getLogger(__name__)
+
+Priority = str
+SEGMENT_RAG = "rag"
+SEGMENT_SESSION = "session"
+SEGMENT_PROFILE = "profile"
+SEGMENT_SYSTEM = "system"
+
+DEFAULT_PRIORITIES: tuple[Priority, ...] = (
+    SEGMENT_SYSTEM,
+    SEGMENT_PROFILE,
+    SEGMENT_SESSION,
+    SEGMENT_RAG,
+)
+
+
+@dataclass
+class BudgetReport:
+    """Observability for a single context build.
+
+    ``used_tokens`` is the final size of the assembled ``system_prompt``
+    counted via the supplied ``TokenCounter``. ``dropped_blocks`` lists
+    block identifiers that did not fit; UI may surface this to the user.
+    """
+
+    used_tokens: int = 0
+    budget: int = 0
+    remaining_tokens: int = 0
+    dropped_blocks: list[str] = field(default_factory=list)
+    section_tokens: dict[str, int] = field(default_factory=dict)
+
+
+@dataclass
+class _Candidate:
+    section: Priority
+    text: str
+    label: str
+
+
+class TokenBudgetedContextBuilder(ContextBuilder):
+    """ContextBuilder that enforces a hard token ceiling on the final
+    ``system_prompt``.
+
+    Behaviour:
+    1. ``system_prompt`` is mandatory and never truncated. If it does not
+       fit into ``max_tokens`` alone, ``TokenBudgetExceededError`` is
+       raised.
+    2. ``profile_text`` is appended in full when ``include_profile`` is
+       true; if it would push us over budget, the profile is dropped and
+       ``dropped_blocks`` is updated (not raised — profile is a hint).
+    3. RAG and session blocks are pooled (top_k * 2) and allocated
+       greedily in priority order. The last accepted block is truncated
+       at a word boundary if it does not fit whole.
+    4. ``BudgetReport`` is attached to the returned ``ContextOutput``
+       via ``budget_report`` so the caller can surface a usage indicator.
+    """
+
+    def __init__(
+        self,
+        store: StoreProtocol,
+        llm: LLMClientProtocol,
+        counter: TokenCounter | None = None,
+        max_tokens: int = 4096,
+        priorities: tuple[Priority, ...] = DEFAULT_PRIORITIES,
+    ) -> None:
+        super().__init__(store, llm)
+        self._counter: TokenCounter = counter or RegexTokenCounter()
+        self._max_tokens = max_tokens
+        self._priorities = priorities
+
+    async def build(self, inp: ContextInput) -> ContextOutput:
+        report = BudgetReport(budget=self._max_tokens)
+
+        system_cost = self._counter.count(inp.system_prompt) if inp.system_prompt else 0
+        if system_cost > self._max_tokens:
+            raise TokenBudgetExceededError(system_cost, self._max_tokens, SEGMENT_SYSTEM)
+        report.section_tokens[SEGMENT_SYSTEM] = system_cost
+        remaining = self._max_tokens - system_cost
+
+        if inp.include_profile and inp.profile_text:
+            profile_block = f"Профиль пользователя:\n{inp.profile_text}"
+            profile_cost = self._counter.count(profile_block)
+            if profile_cost > remaining:
+                logger.warning(
+                    "Profile block (%d tokens) exceeds remaining budget (%d); dropping",
+                    profile_cost, remaining,
+                )
+                report.dropped_blocks.append(SEGMENT_PROFILE)
+                profile_block = ""
+            else:
+                report.section_tokens[SEGMENT_PROFILE] = profile_cost
+                remaining -= profile_cost
+        else:
+            profile_block = ""
+
+        pool: dict[Priority, list[_Candidate]] = {p: [] for p in self._priorities}
+
+        if (inp.include_rag and inp.doc_ids) or (inp.include_session and inp.chat_id):
+            query_emb = (await self._llm.embed([inp.query], model=inp.embedding_model))[0]
+
+            if inp.include_rag and inp.doc_ids:
+                str_doc_ids = [str(d) for d in inp.doc_ids]
+                where = (
+                    {"doc_id": {"$in": str_doc_ids}}
+                    if len(str_doc_ids) > 1
+                    else {"doc_id": str_doc_ids[0]}
+                )
+                rag_hits = self._store.query(
+                    query_emb, top_k=max(1, inp.top_k_rag * 2), where=where
+                )
+                for i, hit in enumerate(rag_hits):
+                    pool[SEGMENT_RAG].append(_Candidate(
+                        section=SEGMENT_RAG,
+                        text=hit["document"],
+                        label=f"rag[{i}]",
+                    ))
+
+            if inp.include_session and inp.chat_id:
+                session_hits = self._store.query(
+                    query_emb,
+                    top_k=max(1, inp.top_k_session * 2),
+                    where={"doc_id": f"session_{inp.chat_id}"},
+                )
+                for i, hit in enumerate(session_hits):
+                    pool[SEGMENT_SESSION].append(_Candidate(
+                        section=SEGMENT_SESSION,
+                        text=hit["document"],
+                        label=f"session[{i}]",
+                    ))
+
+        kept_rag: list[str] = []
+        kept_session: list[str] = []
+
+        for section in self._priorities:
+            if section in (SEGMENT_SYSTEM, SEGMENT_PROFILE):
+                continue
+            if section not in pool or not pool[section]:
+                continue
+            for cand in pool[section]:
+                cost = self._counter.count(cand.text)
+                if cost <= remaining:
+                    if cand.section == SEGMENT_RAG:
+                        kept_rag.append(cand.text)
+                    elif cand.section == SEGMENT_SESSION:
+                        kept_session.append(cand.text)
+                    remaining -= cost
+                    report.section_tokens[cand.label] = cost
+                else:
+                    trimmed = self._truncate_to_budget(cand.text, remaining)
+                    if trimmed:
+                        if cand.section == SEGMENT_RAG:
+                            kept_rag.append(trimmed)
+                        elif cand.section == SEGMENT_SESSION:
+                            kept_session.append(trimmed)
+                        report.section_tokens[cand.label] = remaining
+                        remaining = 0
+                    else:
+                        report.dropped_blocks.append(cand.label)
+                    break
+            if remaining <= 0:
+                for later_section in self._priorities:
+                    if later_section <= section:
+                        continue
+                    for cand in pool.get(later_section, []):
+                        report.dropped_blocks.append(cand.label)
+                break
+
+        report.used_tokens = self._max_tokens - remaining
+        report.remaining_tokens = remaining
+
+        parts: list[str] = []
+        if inp.system_prompt:
+            parts.append(inp.system_prompt)
+        if profile_block:
+            parts.append(profile_block)
+        if kept_rag:
+            parts.append("\n\n---\n\n".join(kept_rag))
+        if kept_session:
+            parts.append("История диалога (сжатая):\n" + "\n---\n".join(kept_session))
+
+        return ContextOutput(
+            system_prompt="\n\n".join(parts),
+            rag_blocks=kept_rag,
+            session_blocks=kept_session,
+            profile_used=bool(profile_block),
+            budget_report=report,
+        )
+
+    def _truncate_to_budget(self, text: str, budget: int) -> str:
+        """Cut ``text`` so it fits into ``budget`` tokens, ending on a
+        word boundary. Returns empty string if no content fits.
+        """
+        if budget <= 0:
+            return ""
+        words = text.split()
+        out: list[str] = []
+        used = 0
+        for w in words:
+            w_cost = self._counter.count(w) + 1
+            if used + w_cost > budget:
+                break
+            out.append(w)
+            used += w_cost
+        if not out:
+            return ""
+        return " ".join(out) + "…"

protoprompt/llm.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class LLMClientProtocol(Protocol):
+    async def chat(self, messages: list[dict], model: str = "", **options: object) -> str:
+        ...
+
+    async def embed(self, texts: list[str], model: str = "") -> list[list[float]]:
+        ...

protoprompt/pipeline.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from protoprompt.llm import LLMClientProtocol
+from protoprompt.session.compressor import Compressor
+from protoprompt.session.strategy import HeuristicStrategy, StrategyProtocol
+from protoprompt.session.types import CompressedBlock, Session
+from protoprompt.store.protocol import StoreProtocol
+
+logger = logging.getLogger(__name__)
+
+
+class Pipeline:
+    """Orchestrate session compression and persistence into a vector store.
+
+    The pipeline decides when compression should run (``should_compress``)
+    and is the only place that should write compressed session data to
+    the store. The write is performed atomically: the new doc_id is
+    written first, then the old one is removed. If the process dies in
+    between, the next call overwrites both, so no chunk is lost.
+    """
+
+    def __init__(
+        self,
+        store: StoreProtocol,
+        llm: LLMClientProtocol,
+        strategy: StrategyProtocol | None = None,
+        compress_every_n: int = 10,
+        embedding_model: str = "nomic-embed-text",
+    ) -> None:
+        self._store = store
+        self._llm = llm
+        self._compressor = Compressor(strategy or HeuristicStrategy())
+        self._compress_every_n = compress_every_n
+        self._embedding_model = embedding_model
+
+    async def compress_and_store(self, session: Session) -> list[CompressedBlock]:
+        if len(session.messages) < self._compress_every_n:
+            return []
+
+        blocks = await self._compressor.compress(session, self._llm)
+        if not blocks:
+            return []
+
+        doc_id = f"session_{session.chat_id}"
+        new_doc_id = f"{doc_id}_new"
+        texts = [b.text for b in blocks]
+        embeddings = await self._llm.embed(texts, model=self._embedding_model)
+
+        meta: dict[str, Any] = {
+            "chat_id": session.chat_id,
+            "strategy": session.strategy,
+            "message_count": len(session.messages),
+        }
+        self._store.add(new_doc_id, texts, embeddings, meta)
+        self._store.delete(doc_id)
+        self._store.add(doc_id, texts, embeddings, meta)
+        self._store.delete(new_doc_id)
+
+        logger.info(
+            "Compressed session %s: %d messages -> %d blocks",
+            session.chat_id,
+            len(session.messages),
+            len(blocks),
+        )
+        return blocks
+
+    def should_compress(self, message_count: int) -> bool:
+        return message_count >= self._compress_every_n

protoprompt/profile/__init__.py

@@ -0,0 +1,4 @@
+from protoprompt.profile.builder import ProfileBuilder
+from protoprompt.profile.types import UserProfile
+
+__all__ = ["ProfileBuilder", "UserProfile"]

protoprompt/profile/builder.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import json
+import logging
+
+from protoprompt.llm import LLMClientProtocol
+from protoprompt.profile.types import UserProfile
+
+logger = logging.getLogger(__name__)
+
+
+class ProfileBuilder:
+    """Build a structured user profile by asking the LLM to analyse turns.
+
+    The LLM is expected to return strict JSON; on any failure the builder
+    logs the original exception and returns a minimal profile with
+    ``summary="Не удалось построить профиль"``.
+    """
+
+    def __init__(self, llm: LLMClientProtocol) -> None:
+        self._llm = llm
+
+    async def build(self, user_id: str, messages: list[dict]) -> UserProfile:
+        if not messages:
+            return UserProfile(user_id=user_id)
+
+        user_texts = [m["content"] for m in messages if m.get("role") == "user"]
+        if not user_texts:
+            return UserProfile(user_id=user_id)
+
+        prompt = (
+            "Проанализируй сообщения пользователя. Выдели:\n"
+            "1. Стиль общения (кратко/развёрнуто, формально/неформально)\n"
+            "2. Предпочтения (любит списки, нарратив, технические детали)\n"
+            "3. Уровень экспертизы (новичок, средний, эксперт)\n"
+            "Ответ дай строго в формате JSON:\n"
+            '{"traits": {"style": "...", "expertise": "..."},'
+            '"preferences": {"format": "..."}, "summary": "..."}\n\n'
+            "Сообщения пользователя:\n" + "\n".join(user_texts[-20:])
+        )
+
+        try:
+            response = await self._llm.chat(
+                [{"role": "user", "content": prompt}],
+                temperature=0.3,
+                max_tokens=300,
+            )
+            data = json.loads(response)
+            return UserProfile(
+                user_id=user_id,
+                traits=data.get("traits", {}),
+                preferences=data.get("preferences", {}),
+                summary=data.get("summary", ""),
+            )
+        except Exception:
+            logger.warning(
+                "Failed to build profile for user_id=%s", user_id, exc_info=True
+            )
+            return UserProfile(
+                user_id=user_id,
+                summary="Не удалось построить профиль",
+            )

protoprompt/profile/schema.py

@@ -0,0 +1,24 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "object",
+  "properties": {
+    "traits": {
+      "type": "object",
+      "properties": {
+        "style": { "type": "string" },
+        "expertise": { "type": "string", "enum": ["beginner", "intermediate", "expert"] },
+        "verbosity": { "type": "string", "enum": ["concise", "balanced", "detailed"] },
+        "formality": { "type": "string", "enum": ["casual", "neutral", "formal"] }
+      }
+    },
+    "preferences": {
+      "type": "object",
+      "properties": {
+        "format": { "type": "string", "enum": ["bullets", "narrative", "code_heavy", "mixed"] },
+        "language": { "type": "string" },
+        "topics": { "type": "array", "items": { "type": "string" } }
+      }
+    },
+    "summary": { "type": "string" }
+  }
+}

protoprompt/profile/types.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class UserProfile:
+    user_id: str = ""
+    traits: dict[str, str] = field(default_factory=dict)
+    preferences: dict[str, str] = field(default_factory=dict)
+    summary: str = ""
+    updated_at: str = ""

protoprompt/session/__init__.py

@@ -0,0 +1,16 @@
+from protoprompt.session.compressor import Compressor
+from protoprompt.session.strategy import (
+    HeuristicStrategy,
+    LLMSummaryStrategy,
+    StrategyProtocol,
+)
+from protoprompt.session.types import CompressedBlock, Session
+
+__all__ = [
+    "Compressor",
+    "HeuristicStrategy",
+    "LLMSummaryStrategy",
+    "StrategyProtocol",
+    "Session",
+    "CompressedBlock",
+]

protoprompt/session/compressor.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import logging
+
+from protoprompt.llm import LLMClientProtocol
+from protoprompt.session.strategy import HeuristicStrategy, StrategyProtocol
+from protoprompt.session.types import CompressedBlock, Session
+
+logger = logging.getLogger(__name__)
+
+
+class Compressor:
+    def __init__(self, strategy: StrategyProtocol | None = None) -> None:
+        self._strategy = strategy or HeuristicStrategy()
+
+    async def compress(self, session: Session, llm: LLMClientProtocol) -> list[CompressedBlock]:
+        return await self._strategy.compress(session, llm)