loopshield 1.0.0__tar.gz

loopshield-1.0.0/PKG-INFO

@@ -0,0 +1,229 @@
+Metadata-Version: 2.4
+Name: loopshield
+Version: 1.0.0
+Summary: A lightweight, vendor-agnostic AI agent loop detection and self-healing middleware
+Author-email: Your Name <your.email@example.com>
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Framework :: FastAPI
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: networkx>=3.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: pydantic>=2.0
+Requires-Dist: fastapi>=0.100.0
+Requires-Dist: uvicorn>=0.22.0
+Requires-Dist: httpx>=0.24.0
+Provides-Extra: langchain
+Requires-Dist: langchain>=0.1.0; extra == "langchain"
+Requires-Dist: langgraph>=0.1.0; extra == "langchain"
+Provides-Extra: langgraph
+Requires-Dist: langgraph>=0.1.0; extra == "langgraph"
+Provides-Extra: crewai
+Requires-Dist: crewai>=0.1.0; extra == "crewai"
+Provides-Extra: all
+Requires-Dist: loopshield[crewai,langchain,langgraph]; extra == "all"
+
+# LoopShield
+
+LoopShield is a vendor-agnostic Python SDK and an OpenAI-compatible proxy that stops AI agents from burning tokens on the same tool cycle or the same thought. The proxy sits in front of the model provider. Every `POST /v1/chat/completions` is screened before it is forwarded. A detected loop comes back as a normal assistant message that tells the agent the repeating tool is locked and it must change strategy. The upstream provider is not called.
+
+The core package depends only on `networkx`, `numpy`, and Pydantic v2. FastAPI, uvicorn, and httpx are an optional `proxy` extra. Embedding models are injected; the OpenAI and Ollama SDKs are not core dependencies.
+
+## How a call is intercepted
+
+Clients keep speaking the provider’s chat API. LoopShield accepts an OpenAI Chat Completions body, and it also accepts an Anthropic Messages body on the same route. Anthropic shape is recognized from an `anthropic-version` header, a payload that has both `system` and `max_tokens`, or content blocks of type `tool_use` / `tool_result`.
+
+```mermaid
+flowchart TD
+    agent["Agent or app"]
+    route["POST /v1/chat/completions"]
+    style{"Payload style"}
+    steps["Assistant turns become AgentSteps"]
+    inspect["LoopShield.inspect"]
+    graph["GraphAnalyzer<br/>exact tool-call cycle"]
+    semantic["SemanticAnalyzer<br/>cosine similarity of thoughts"]
+    loop{"Loop?"}
+    nudge["Intercepted assistant message<br/>intervention_prompt"]
+    budget{"Over SHIELD_MAX_TOKENS?"}
+    prune["prune_context<br/>keep system prompt and last 2 turns"]
+    openai["Upstream POST /v1/chat/completions"]
+    anthropic["Upstream POST /v1/messages"]
+
+    agent --> route --> style --> steps --> inspect
+    inspect --> graph
+    inspect --> semantic
+    graph --> loop
+    semantic --> loop
+    loop -->|yes| nudge
+    nudge --> agent
+    loop -->|no| budget
+    budget -->|yes| prune --> forward{"Provider"}
+    budget -->|no| forward
+    forward -->|OpenAI| openai
+    forward -->|Anthropic| anthropic
+    openai --> agent
+    anthropic --> agent
+```
+
+A clean request is forwarded with the caller’s `Authorization`, `x-api-key`, `anthropic-version`, `anthropic-beta`, `openai-organization`, and `openai-project` headers. LoopShield does not store those credentials.
+
+When a loop is found, the JSON response still looks like a completion:
+
+- OpenAI: `object` is `chat.completion`, and `choices[0].message.content` is the nudge.
+- Anthropic: `type` is `message`, `role` is `assistant`, and the text block is the nudge.
+- Both include `loopshield.intercepted`, `loopshield.loop_type` (`graph_cycle`, `semantic`, or `both`), and `loopshield.intervention_prompt`.
+
+`stream: true` on an intercepted call is returned as server-sent events with the same nudge. `GET /health` returns `{"status": "ok"}`.
+
+## Quick start
+
+Build and run the gateway on port 8000. Point `UPSTREAM_API_BASE` at the provider you want clean calls to reach.
+
+```bash
+docker build -t loopshield:0.1.0 .
+docker run --rm -p 8000:8000 -e UPSTREAM_API_BASE=https://api.openai.com -e SHIELD_SIMILARITY_THRESHOLD=0.92 loopshield:0.1.0
+```
+
+Check the process, then send a chat completion through it. The provider key stays on the request; the container only forwards it.
+
+```bash
+curl http://127.0.0.1:8000/health
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H "Authorization: Bearer $OPENAI_API_KEY" \
+  -H "Content-Type: application/json" \
+  -d "{\"model\":\"gpt-4.1\",\"messages\":[{\"role\":\"user\",\"content\":\"Summarize the deploy log.\"}]}"
+```
+
+An OpenAI SDK client uses the proxy as its base URL:
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="sk-...")
+completion = client.chat.completions.create(
+    model="gpt-4.1",
+    messages=[{"role": "user", "content": "Summarize the deploy log."}],
+)
+```
+
+Anthropic-shaped traffic uses the same path. Set the upstream base to `https://api.anthropic.com` and send the provider headers the proxy already forwards:
+
+```bash
+docker run --rm -p 8000:8000 -e UPSTREAM_API_BASE=https://api.anthropic.com loopshield:0.1.0
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H "x-api-key: $ANTHROPIC_API_KEY" \
+  -H "anthropic-version: 2023-06-01" \
+  -H "Content-Type: application/json" \
+  -d "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":256,\"system\":\"Be careful.\",\"messages\":[{\"role\":\"user\",\"content\":\"Summarize the deploy log.\"}]}"
+```
+
+The official Anthropic SDK posts to `/v1/messages`. This gateway accepts that payload shape on `/v1/chat/completions`, then forwards a clean call to `{UPSTREAM_API_BASE}/v1/messages`. If the upstream base already ends in `/v1`, that prefix is not doubled.
+
+## Configuration
+
+| Variable | Default | Meaning |
+| --- | --- | --- |
+| `UPSTREAM_API_BASE` | `https://api.openai.com` | Origin for forwarded calls. OpenAI traffic goes to `/v1/chat/completions`. Anthropic traffic goes to `/v1/messages`. |
+| `SHIELD_SIMILARITY_THRESHOLD` | `0.92` | Cosine similarity, from 0 to 1, above which two thoughts count as the same line of reasoning. The comparison is strict `>`. |
+| `SHIELD_MAX_TOKENS` | `8000` | Estimated token budget. Over this, `prune_context` keeps the leading system prompt and the last two turns, and collapses repeated tool errors in the middle. |
+| `SHIELD_MAX_CYCLES` | `1` | A directed tool cycle is a loop when it completes more times than this. `1` allows one lap and flags the second. `0` flags the first time the cycle closes. |
+| `SHIELD_MAX_SEMANTIC_REPEATS` | `2` | Minimum streak of consecutive similar thoughts, measured backward from the latest thought, required to flag a semantic loop. |
+| `SHIELD_UPSTREAM_TIMEOUT` | `60` | Upstream HTTP timeout in seconds. |
+
+The image sets `UPSTREAM_API_BASE`, `SHIELD_SIMILARITY_THRESHOLD`, and `SHIELD_MAX_TOKENS`. Override any of them with `docker run -e`.
+
+Two detectors run on every request:
+
+- **Graph.** Each tool call is a node of `tool_name` plus normalized arguments. Consecutive calls are directed edges. The same cycle walked again, including a tool calling itself with the same arguments, is a graph loop.
+- **Semantic.** Thoughts are compared with cosine similarity. With no embedding model configured, the proxy uses an in-process TF-IDF bag-of-words embedder. Identical text scores 1.0. A streak breaks at the first thought that does not exceed the threshold.
+
+The proxy calls `inspect`, which does not keep request history on the shared process. Concurrent requests do not contaminate each other.
+
+## Use the SDK in process
+
+Install the library when the agent loop lives in your own code and you do not need the gateway:
+
+```bash
+pip install -e ".[dev]"
+```
+
+Record each step. A verdict with `is_loop` carries the system nudge to send back to the model.
+
+```python
+from loopshield import AgentStep, LoopShield
+
+shield = LoopShield(similarity_threshold=0.92, max_cycles=1)
+
+verdict = shield.step(
+    AgentStep(
+        thought="The config file should be at the project root.",
+        tool_name="read_file",
+        tool_args={"path": "config.txt"},
+        tokens_used=12,
+    )
+)
+
+if verdict.is_loop:
+    print(verdict.intervention_prompt)
+```
+
+`@shield.guard` wraps a function that takes an `AgentStep`. When the new step closes a loop, the decorator returns the `LoopVerdict` and does not call the function. Pass an initialized OpenAI or Ollama client as `embedding=` to replace the TF-IDF fallback; pass `embed_fn=` for any other model. Do not pass both.
+
+`prune_context(messages, max_tokens_estimate)` takes OpenAI-style message dicts. It estimates tokens from `token_count`, `tokens`, or `usage.total_tokens` when present, otherwise from character length at about four characters per token.
+
+## Framework integrations
+
+`loopshield/frameworks.py` plugs the same engine into LangChain, CrewAI, and LangGraph. Each framework is an optional install (`pip install "loopshield[langchain]"`, `[langgraph]`, `[crewai]`). The hooks keep no history on the shield, so one `LoopShield` can serve concurrent runs.
+
+**LangChain.** `shield.langchain_middleware(config)` returns an `AgentMiddleware`:
+
+```python
+from langchain.agents import create_agent
+
+agent = create_agent(model, tools, middleware=[shield.langchain_middleware({"max_tokens": 6000})])
+```
+
+- `wrap_tool_call` replaces a looping tool call with an error `ToolMessage` that carries the nudge. The tool does not run, and the model reads why on its next turn.
+- `wrap_model_call` runs `prune_context` when the transcript is over budget. For a loop that reaches the model, it answers with the nudge before the LLM is called. Set `on_loop` to `inject` to append the nudge as a system message and let the model run.
+
+**CrewAI.** `@shield.crewai_agent()` wraps `Agent.execute_task`, either on a class or on the method. It rebuilds the tool trajectory from `agent.tools_results`, or from the `Action` / `Action Input` lines in the string context. When the tools loop, the nudge is prepended to the `context` string that the next step reads.
+
+```python
+@shield.crewai_agent()
+class GuardedAgent(Agent): ...
+```
+
+**LangGraph.** `shield.langgraph_node()` wraps a node. A looping `state["messages"]` gets the system nudge as a state update, and the node is skipped for that pass. With `{"retry_node": "agent"}` it returns `Command(update=..., goto="agent")`, so the graph re-enters that node with the nudge in context. The `messages` channel needs an appending reducer such as `add_messages`.
+
+```python
+graph.add_node("agent", shield.langgraph_node(agent_node, config={"retry_node": "agent"}))
+```
+
+## Development
+
+```bash
+python -m pytest
+```
+
+Python 3.11 or newer. Tests live in `tests/` and cover tool-call cycles, semantic streaks, recovery after a nudge, context pruning, embedding adapters, and the proxy with a mocked upstream.
+
+| Module | Role |
+| --- | --- |
+| `loopshield/models.py` | `AgentStep` and `LoopVerdict` |
+| `loopshield/graph_analyzer.py` | Directed tool-call cycles |
+| `loopshield/semantic_analyzer.py` | Cosine similarity of thoughts |
+| `loopshield/mitigator.py` | Intervention nudge and context pruning |
+| `loopshield/shield.py` | `LoopShield.step`, `inspect`, and `guard` |
+| `loopshield/adapters.py` | Optional OpenAI and Ollama embedding wrappers |
+| `loopshield/proxy.py` | FastAPI gateway served by uvicorn |
+| `loopshield/transcript.py` | Chat messages from any framework to `AgentStep`s |
+| `loopshield/frameworks.py` | LangChain, CrewAI, and LangGraph integrations |

loopshield-1.0.0/README.md

@@ -0,0 +1,201 @@
+# LoopShield
+
+LoopShield is a vendor-agnostic Python SDK and an OpenAI-compatible proxy that stops AI agents from burning tokens on the same tool cycle or the same thought. The proxy sits in front of the model provider. Every `POST /v1/chat/completions` is screened before it is forwarded. A detected loop comes back as a normal assistant message that tells the agent the repeating tool is locked and it must change strategy. The upstream provider is not called.
+
+The core package depends only on `networkx`, `numpy`, and Pydantic v2. FastAPI, uvicorn, and httpx are an optional `proxy` extra. Embedding models are injected; the OpenAI and Ollama SDKs are not core dependencies.
+
+## How a call is intercepted
+
+Clients keep speaking the provider’s chat API. LoopShield accepts an OpenAI Chat Completions body, and it also accepts an Anthropic Messages body on the same route. Anthropic shape is recognized from an `anthropic-version` header, a payload that has both `system` and `max_tokens`, or content blocks of type `tool_use` / `tool_result`.
+
+```mermaid
+flowchart TD
+    agent["Agent or app"]
+    route["POST /v1/chat/completions"]
+    style{"Payload style"}
+    steps["Assistant turns become AgentSteps"]
+    inspect["LoopShield.inspect"]
+    graph["GraphAnalyzer<br/>exact tool-call cycle"]
+    semantic["SemanticAnalyzer<br/>cosine similarity of thoughts"]
+    loop{"Loop?"}
+    nudge["Intercepted assistant message<br/>intervention_prompt"]
+    budget{"Over SHIELD_MAX_TOKENS?"}
+    prune["prune_context<br/>keep system prompt and last 2 turns"]
+    openai["Upstream POST /v1/chat/completions"]
+    anthropic["Upstream POST /v1/messages"]
+
+    agent --> route --> style --> steps --> inspect
+    inspect --> graph
+    inspect --> semantic
+    graph --> loop
+    semantic --> loop
+    loop -->|yes| nudge
+    nudge --> agent
+    loop -->|no| budget
+    budget -->|yes| prune --> forward{"Provider"}
+    budget -->|no| forward
+    forward -->|OpenAI| openai
+    forward -->|Anthropic| anthropic
+    openai --> agent
+    anthropic --> agent
+```
+
+A clean request is forwarded with the caller’s `Authorization`, `x-api-key`, `anthropic-version`, `anthropic-beta`, `openai-organization`, and `openai-project` headers. LoopShield does not store those credentials.
+
+When a loop is found, the JSON response still looks like a completion:
+
+- OpenAI: `object` is `chat.completion`, and `choices[0].message.content` is the nudge.
+- Anthropic: `type` is `message`, `role` is `assistant`, and the text block is the nudge.
+- Both include `loopshield.intercepted`, `loopshield.loop_type` (`graph_cycle`, `semantic`, or `both`), and `loopshield.intervention_prompt`.
+
+`stream: true` on an intercepted call is returned as server-sent events with the same nudge. `GET /health` returns `{"status": "ok"}`.
+
+## Quick start
+
+Build and run the gateway on port 8000. Point `UPSTREAM_API_BASE` at the provider you want clean calls to reach.
+
+```bash
+docker build -t loopshield:0.1.0 .
+docker run --rm -p 8000:8000 -e UPSTREAM_API_BASE=https://api.openai.com -e SHIELD_SIMILARITY_THRESHOLD=0.92 loopshield:0.1.0
+```
+
+Check the process, then send a chat completion through it. The provider key stays on the request; the container only forwards it.
+
+```bash
+curl http://127.0.0.1:8000/health
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H "Authorization: Bearer $OPENAI_API_KEY" \
+  -H "Content-Type: application/json" \
+  -d "{\"model\":\"gpt-4.1\",\"messages\":[{\"role\":\"user\",\"content\":\"Summarize the deploy log.\"}]}"
+```
+
+An OpenAI SDK client uses the proxy as its base URL:
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="sk-...")
+completion = client.chat.completions.create(
+    model="gpt-4.1",
+    messages=[{"role": "user", "content": "Summarize the deploy log."}],
+)
+```
+
+Anthropic-shaped traffic uses the same path. Set the upstream base to `https://api.anthropic.com` and send the provider headers the proxy already forwards:
+
+```bash
+docker run --rm -p 8000:8000 -e UPSTREAM_API_BASE=https://api.anthropic.com loopshield:0.1.0
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H "x-api-key: $ANTHROPIC_API_KEY" \
+  -H "anthropic-version: 2023-06-01" \
+  -H "Content-Type: application/json" \
+  -d "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":256,\"system\":\"Be careful.\",\"messages\":[{\"role\":\"user\",\"content\":\"Summarize the deploy log.\"}]}"
+```
+
+The official Anthropic SDK posts to `/v1/messages`. This gateway accepts that payload shape on `/v1/chat/completions`, then forwards a clean call to `{UPSTREAM_API_BASE}/v1/messages`. If the upstream base already ends in `/v1`, that prefix is not doubled.
+
+## Configuration
+
+| Variable | Default | Meaning |
+| --- | --- | --- |
+| `UPSTREAM_API_BASE` | `https://api.openai.com` | Origin for forwarded calls. OpenAI traffic goes to `/v1/chat/completions`. Anthropic traffic goes to `/v1/messages`. |
+| `SHIELD_SIMILARITY_THRESHOLD` | `0.92` | Cosine similarity, from 0 to 1, above which two thoughts count as the same line of reasoning. The comparison is strict `>`. |
+| `SHIELD_MAX_TOKENS` | `8000` | Estimated token budget. Over this, `prune_context` keeps the leading system prompt and the last two turns, and collapses repeated tool errors in the middle. |
+| `SHIELD_MAX_CYCLES` | `1` | A directed tool cycle is a loop when it completes more times than this. `1` allows one lap and flags the second. `0` flags the first time the cycle closes. |
+| `SHIELD_MAX_SEMANTIC_REPEATS` | `2` | Minimum streak of consecutive similar thoughts, measured backward from the latest thought, required to flag a semantic loop. |
+| `SHIELD_UPSTREAM_TIMEOUT` | `60` | Upstream HTTP timeout in seconds. |
+
+The image sets `UPSTREAM_API_BASE`, `SHIELD_SIMILARITY_THRESHOLD`, and `SHIELD_MAX_TOKENS`. Override any of them with `docker run -e`.
+
+Two detectors run on every request:
+
+- **Graph.** Each tool call is a node of `tool_name` plus normalized arguments. Consecutive calls are directed edges. The same cycle walked again, including a tool calling itself with the same arguments, is a graph loop.
+- **Semantic.** Thoughts are compared with cosine similarity. With no embedding model configured, the proxy uses an in-process TF-IDF bag-of-words embedder. Identical text scores 1.0. A streak breaks at the first thought that does not exceed the threshold.
+
+The proxy calls `inspect`, which does not keep request history on the shared process. Concurrent requests do not contaminate each other.
+
+## Use the SDK in process
+
+Install the library when the agent loop lives in your own code and you do not need the gateway:
+
+```bash
+pip install -e ".[dev]"
+```
+
+Record each step. A verdict with `is_loop` carries the system nudge to send back to the model.
+
+```python
+from loopshield import AgentStep, LoopShield
+
+shield = LoopShield(similarity_threshold=0.92, max_cycles=1)
+
+verdict = shield.step(
+    AgentStep(
+        thought="The config file should be at the project root.",
+        tool_name="read_file",
+        tool_args={"path": "config.txt"},
+        tokens_used=12,
+    )
+)
+
+if verdict.is_loop:
+    print(verdict.intervention_prompt)
+```
+
+`@shield.guard` wraps a function that takes an `AgentStep`. When the new step closes a loop, the decorator returns the `LoopVerdict` and does not call the function. Pass an initialized OpenAI or Ollama client as `embedding=` to replace the TF-IDF fallback; pass `embed_fn=` for any other model. Do not pass both.
+
+`prune_context(messages, max_tokens_estimate)` takes OpenAI-style message dicts. It estimates tokens from `token_count`, `tokens`, or `usage.total_tokens` when present, otherwise from character length at about four characters per token.
+
+## Framework integrations
+
+`loopshield/frameworks.py` plugs the same engine into LangChain, CrewAI, and LangGraph. Each framework is an optional install (`pip install "loopshield[langchain]"`, `[langgraph]`, `[crewai]`). The hooks keep no history on the shield, so one `LoopShield` can serve concurrent runs.
+
+**LangChain.** `shield.langchain_middleware(config)` returns an `AgentMiddleware`:
+
+```python
+from langchain.agents import create_agent
+
+agent = create_agent(model, tools, middleware=[shield.langchain_middleware({"max_tokens": 6000})])
+```
+
+- `wrap_tool_call` replaces a looping tool call with an error `ToolMessage` that carries the nudge. The tool does not run, and the model reads why on its next turn.
+- `wrap_model_call` runs `prune_context` when the transcript is over budget. For a loop that reaches the model, it answers with the nudge before the LLM is called. Set `on_loop` to `inject` to append the nudge as a system message and let the model run.
+
+**CrewAI.** `@shield.crewai_agent()` wraps `Agent.execute_task`, either on a class or on the method. It rebuilds the tool trajectory from `agent.tools_results`, or from the `Action` / `Action Input` lines in the string context. When the tools loop, the nudge is prepended to the `context` string that the next step reads.
+
+```python
+@shield.crewai_agent()
+class GuardedAgent(Agent): ...
+```
+
+**LangGraph.** `shield.langgraph_node()` wraps a node. A looping `state["messages"]` gets the system nudge as a state update, and the node is skipped for that pass. With `{"retry_node": "agent"}` it returns `Command(update=..., goto="agent")`, so the graph re-enters that node with the nudge in context. The `messages` channel needs an appending reducer such as `add_messages`.
+
+```python
+graph.add_node("agent", shield.langgraph_node(agent_node, config={"retry_node": "agent"}))
+```
+
+## Development
+
+```bash
+python -m pytest
+```
+
+Python 3.11 or newer. Tests live in `tests/` and cover tool-call cycles, semantic streaks, recovery after a nudge, context pruning, embedding adapters, and the proxy with a mocked upstream.
+
+| Module | Role |
+| --- | --- |
+| `loopshield/models.py` | `AgentStep` and `LoopVerdict` |
+| `loopshield/graph_analyzer.py` | Directed tool-call cycles |
+| `loopshield/semantic_analyzer.py` | Cosine similarity of thoughts |
+| `loopshield/mitigator.py` | Intervention nudge and context pruning |
+| `loopshield/shield.py` | `LoopShield.step`, `inspect`, and `guard` |
+| `loopshield/adapters.py` | Optional OpenAI and Ollama embedding wrappers |
+| `loopshield/proxy.py` | FastAPI gateway served by uvicorn |
+| `loopshield/transcript.py` | Chat messages from any framework to `AgentStep`s |
+| `loopshield/frameworks.py` | LangChain, CrewAI, and LangGraph integrations |

loopshield-1.0.0/loopshield/__init__.py

@@ -0,0 +1,25 @@
+"""LoopShield: loop detection and self-healing middleware for AI agents."""
+
+from loopshield.adapters import BaseEmbeddingAdapter, OllamaEmbeddingAdapter, OpenAIEmbeddingAdapter
+from loopshield.graph_analyzer import GraphAnalyzer
+from loopshield.mitigator import Mitigator, generate_nudge, prune_context
+from loopshield.models import AgentStep, LoopVerdict
+from loopshield.semantic_analyzer import SemanticAnalyzer
+from loopshield.shield import LoopShield
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AgentStep",
+    "BaseEmbeddingAdapter",
+    "GraphAnalyzer",
+    "LoopShield",
+    "LoopVerdict",
+    "Mitigator",
+    "OllamaEmbeddingAdapter",
+    "OpenAIEmbeddingAdapter",
+    "SemanticAnalyzer",
+    "generate_nudge",
+    "prune_context",
+    "__version__",
+]

loopshield-1.0.0/loopshield/adapters.py

@@ -0,0 +1,222 @@
+"""Embedding adapters for clients the application already constructs.
+
+The core package does not import OpenAI or Ollama. Pass an initialized client
+and these wrappers call ``client.embeddings.create`` or the local Ollama
+embedding API.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable, Mapping
+from typing import Any, Protocol
+
+
+DEFAULT_OPENAI_EMBEDDING_MODEL = "text-embedding-3-small"
+DEFAULT_OLLAMA_EMBEDDING_MODEL = "nomic-embed-text"
+
+
+class BaseEmbeddingAdapter(Protocol):
+    """Any embedding backend that maps one string to a dense vector."""
+
+    def embed(self, text: str) -> list[float]:
+        """Return a dense vector for ``text``."""
+
+
+class OpenAIEmbeddingsResource(Protocol):
+    """The ``client.embeddings`` resource on an ``openai.OpenAI`` client."""
+
+    def create(self, *, model: str, input: str, **kwargs: Any) -> Any:
+        """Create an embedding for ``input``."""
+
+
+class OpenAIClient(Protocol):
+    """Initialized OpenAI client with an embeddings resource."""
+
+    embeddings: OpenAIEmbeddingsResource
+
+
+class OllamaEmbedClient(Protocol):
+    """Ollama client that exposes the current ``embed`` call."""
+
+    def embed(self, *, model: str, input: str, **kwargs: Any) -> Any:
+        """Embed ``input`` with ``model``."""
+
+
+class OllamaLegacyClient(Protocol):
+    """Ollama client that still exposes the legacy ``embeddings`` call."""
+
+    def embeddings(self, *, model: str, prompt: str, **kwargs: Any) -> Any:
+        """Embed ``prompt`` with ``model``."""
+
+
+class OpenAIEmbeddingAdapter:
+    """Call ``client.embeddings.create`` and return the first vector."""
+
+    def __init__(
+        self,
+        client: OpenAIClient,
+        *,
+        model: str = DEFAULT_OPENAI_EMBEDDING_MODEL,
+    ) -> None:
+        self._client = client
+        self.model = _require_model(model)
+
+    def embed(self, text: str) -> list[float]:
+        """Embed ``text`` with the configured OpenAI model."""
+
+        response = self._client.embeddings.create(model=self.model, input=text)
+        return openai_embedding_vector(response)
+
+    def __call__(self, text: str) -> list[float]:
+        """Allow the adapter to be used anywhere an ``embed_fn`` is accepted."""
+
+        return self.embed(text)
+
+
+class OllamaEmbeddingAdapter:
+    """Call a local Ollama client and return the first vector.
+
+    Prefers ``client.embed(model=..., input=...)``. Falls back to the legacy
+    ``client.embeddings(model=..., prompt=...)`` call when that is all the
+    client provides.
+    """
+
+    def __init__(
+        self,
+        client: OllamaEmbedClient | OllamaLegacyClient,
+        *,
+        model: str = DEFAULT_OLLAMA_EMBEDDING_MODEL,
+    ) -> None:
+        self._client = client
+        self.model = _require_model(model)
+
+    def embed(self, text: str) -> list[float]:
+        """Embed ``text`` with the configured local model."""
+
+        embed = getattr(self._client, "embed", None)
+        if callable(embed):
+            return ollama_embedding_vector(embed(model=self.model, input=text))
+        embeddings = getattr(self._client, "embeddings", None)
+        if callable(embeddings):
+            return ollama_embedding_vector(embeddings(model=self.model, prompt=text))
+        raise TypeError("Ollama client must provide embed() or embeddings().")
+
+    def __call__(self, text: str) -> list[float]:
+        """Allow the adapter to be used anywhere an ``embed_fn`` is accepted."""
+
+        return self.embed(text)
+
+
+def as_embed_fn(source: Any, *, model: str | None = None) -> Callable[[str], list[float]]:
+    """Turn an adapter, embed callable, or OpenAI/Ollama client into an embed function.
+
+    ``model`` is used only when ``source`` is a raw client. Adapters already
+    carry their model.
+    """
+
+    if _is_text_embedding_adapter(source):
+        if model is not None:
+            raise ValueError("Set the model on the adapter, or pass the client with embedding_model.")
+        if callable(source):
+            return source
+        return source.embed
+    if _is_openai_client(source):
+        return OpenAIEmbeddingAdapter(source, model=model or DEFAULT_OPENAI_EMBEDDING_MODEL)
+    if _is_ollama_client(source):
+        return OllamaEmbeddingAdapter(source, model=model or DEFAULT_OLLAMA_EMBEDDING_MODEL)
+    if callable(source):
+        if model is not None:
+            raise ValueError("embedding_model applies to an OpenAI or Ollama client, not a plain embed function.")
+        return source
+    raise TypeError(
+        "embedding must be an embed callable, an embedding adapter, or an initialized OpenAI or Ollama client."
+    )
+
+
+def openai_embedding_vector(response: Any) -> list[float]:
+    """Read the first vector from an OpenAI embeddings response."""
+
+    data = response["data"] if isinstance(response, Mapping) else getattr(response, "data", None)
+    if not data:
+        raise ValueError("OpenAI embedding response did not include any vectors.")
+    first = data[0]
+    embedding = first["embedding"] if isinstance(first, Mapping) else getattr(first, "embedding", None)
+    return _as_floats(embedding)
+
+
+def ollama_embedding_vector(response: Any) -> list[float]:
+    """Read the first vector from a current or legacy Ollama embeddings response."""
+
+    if isinstance(response, Mapping):
+        if "embeddings" in response:
+            return _first_vector(response["embeddings"], provider="Ollama")
+        if "embedding" in response:
+            return _as_floats(response["embedding"])
+    embeddings = getattr(response, "embeddings", None)
+    if embeddings:
+        return _first_vector(embeddings, provider="Ollama")
+    embedding = getattr(response, "embedding", None)
+    if embedding:
+        return _as_floats(embedding)
+    raise ValueError("Ollama embedding response did not include a vector.")
+
+
+def _require_model(model: str) -> str:
+    cleaned = model.strip()
+    if not cleaned:
+        raise ValueError("model must be a non-empty string")
+    return cleaned
+
+
+def _as_floats(values: Any) -> list[float]:
+    if values is None:
+        raise ValueError("Embedding response did not include a vector.")
+    try:
+        vector = [float(item) for item in values]
+    except TypeError as exc:
+        raise ValueError("Embedding response did not contain a numeric vector.") from exc
+    if not vector:
+        raise ValueError("Embedding response was empty.")
+    return vector
+
+
+def _first_vector(vectors: Any, *, provider: str) -> list[float]:
+    try:
+        first = vectors[0]
+    except (IndexError, TypeError, KeyError) as exc:
+        raise ValueError(f"{provider} embedding response did not include any vectors.") from exc
+    return _as_floats(first)
+
+
+def _is_text_embedding_adapter(source: object) -> bool:
+    embed = getattr(source, "embed", None)
+    if not callable(embed):
+        return False
+    try:
+        signature = inspect.signature(embed)
+    except (TypeError, ValueError):
+        return False
+    parameters = [parameter for name, parameter in signature.parameters.items() if name != "self"]
+    if not parameters:
+        return False
+    first = parameters[0]
+    return first.name == "text" and first.kind in (
+        inspect.Parameter.POSITIONAL_ONLY,
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+    )
+
+
+def _is_openai_client(source: object) -> bool:
+    embeddings = getattr(source, "embeddings", None)
+    if embeddings is None or callable(embeddings):
+        return False
+    return callable(getattr(embeddings, "create", None))
+
+
+def _is_ollama_client(source: object) -> bool:
+    if _is_openai_client(source):
+        return False
+    if callable(getattr(source, "embed", None)):
+        return True
+    return callable(getattr(source, "embeddings", None))