add tei support
This commit is contained in:
Nicolò Boschi
parent
f7789f4961
commit
e722a48b14
13 changed files with 729 additions and 447 deletions
16
.env.example
16
.env.example
|
|
@ -13,3 +13,19 @@ HINDSIGHT_API_LOG_LEVEL=info
|
|||
|
||||
# Database (Optional - uses embedded pg0 by default)
|
||||
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
|
||||
|
||||
# Reranker Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
# HINDSIGHT_API_RERANKER_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
|
|
|
|||
147
AGENTS.md
Normal file
147
AGENTS.md
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
# AGENTS.md
|
||||
|
||||
This document captures architectural decisions and coding conventions for the Hindsight project.
|
||||
|
||||
## Documentation
|
||||
|
||||
- **Main documentation**: [hindsight-docs/docs/developer/](./hindsight-docs/docs/developer/)
|
||||
- **Use case patterns**: [hindsight-docs/docs/cookbook/](./hindsight-docs/docs/cookbook/)
|
||||
- **API reference**: Auto-generated from OpenAPI spec
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
hindsight/ # Python package for embedded usage
|
||||
hindsight-api/ # FastAPI server (core memory engine)
|
||||
hindsight-cli/ # Rust CLI client
|
||||
hindsight-control-plane/ # Next.js admin UI
|
||||
hindsight-docs/ # Docusaurus documentation site
|
||||
hindsight-dev/ # Development tools and benchmarks
|
||||
hindsight-integrations/ # Framework integrations (LangChain, etc.)
|
||||
hindsight-clients/ # Generated API clients (Python, TypeScript, Rust)
|
||||
```
|
||||
|
||||
## Core Concepts
|
||||
|
||||
### Memory Banks
|
||||
- Each bank is an isolated memory store (like a "brain" for one user/agent)
|
||||
- Banks contain: memory units (facts), entities, documents, entity links
|
||||
- Banks have a **disposition** (personality traits) and **background** (context)
|
||||
- Bank isolation is strict - no cross-bank data leakage
|
||||
|
||||
### Memory Types
|
||||
- **World facts**: General knowledge ("The sky is blue")
|
||||
- **Experience facts**: Personal experiences ("I visited Paris in 2023")
|
||||
- **Opinion facts**: Beliefs with confidence scores ("Paris is beautiful" - 0.9 confidence)
|
||||
|
||||
### Operations
|
||||
- **Retain**: Store new memories (extracts facts, entities, relationships)
|
||||
- **Recall**: Retrieve memories (semantic, BM25, graph, temporal search)
|
||||
- **Reflect**: Deep analysis to form new insights/opinions
|
||||
|
||||
## API Design Decisions
|
||||
|
||||
### Single Bank Per Request
|
||||
- All API endpoints (`recall`, `reflect`, `retain`) operate on a single bank
|
||||
- Multi-bank queries are the **client/agent's responsibility** to orchestrate
|
||||
- This keeps the API simple and the isolation model clear
|
||||
|
||||
### Disposition Traits (3-trait system)
|
||||
- **Skepticism** (1-5): How skeptical vs trusting when forming opinions
|
||||
- **Literalism** (1-5): How literally to interpret information
|
||||
- **Empathy** (1-5): How much to consider emotional context
|
||||
- These influence the `reflect` operation, not `recall`
|
||||
- Background info also only affects `reflect` (opinion formation)
|
||||
|
||||
## Multi-Bank Architecture Patterns
|
||||
|
||||
See [hindsight-docs/docs/cookbook/](./hindsight-docs/docs/cookbook/) for detailed guides:
|
||||
|
||||
- **Per-User Memory**: One bank per user, simplest pattern
|
||||
- **Support Agent + Shared Knowledge**: User bank + shared docs bank, client orchestrates
|
||||
|
||||
## Developer Guide
|
||||
|
||||
### Running the API Server
|
||||
|
||||
```bash
|
||||
# From project root
|
||||
./scripts/dev/start-api.sh
|
||||
|
||||
# With options
|
||||
./scripts/dev/start-api.sh --reload --port 8888 --log-level debug
|
||||
```
|
||||
|
||||
### Running Tests
|
||||
|
||||
```bash
|
||||
# API tests
|
||||
cd hindsight-api
|
||||
uv run pytest tests/
|
||||
|
||||
# Specific test
|
||||
uv run pytest tests/test_http_api_integration.py -v
|
||||
```
|
||||
|
||||
### Generating OpenAPI Spec
|
||||
|
||||
After changing API endpoints, regenerate the OpenAPI spec and docs:
|
||||
|
||||
```bash
|
||||
./scripts/generate-openapi.sh
|
||||
```
|
||||
|
||||
This will:
|
||||
1. Generate `openapi.json` at project root
|
||||
2. Copy to `hindsight-docs/openapi.json`
|
||||
3. Regenerate API reference documentation
|
||||
|
||||
### Generating API Clients
|
||||
|
||||
After updating the OpenAPI spec, regenerate all clients:
|
||||
|
||||
```bash
|
||||
./scripts/generate-clients.sh
|
||||
```
|
||||
|
||||
This generates:
|
||||
- **Rust client**: `hindsight-clients/rust/` (via progenitor in build.rs)
|
||||
- **Python client**: `hindsight-clients/python/` (via openapi-generator Docker)
|
||||
- **TypeScript client**: `hindsight-clients/typescript/` (via @hey-api/openapi-ts)
|
||||
|
||||
Note: The maintained wrapper `hindsight_client.py` and `README.md` are preserved during regeneration.
|
||||
|
||||
### Running the Documentation Site
|
||||
|
||||
```bash
|
||||
./scripts/dev/start-docs.sh
|
||||
```
|
||||
|
||||
### Running the Control Plane
|
||||
|
||||
```bash
|
||||
./scripts/dev/start-control-plane.sh
|
||||
```
|
||||
|
||||
## Code Style
|
||||
|
||||
### Python (hindsight-api)
|
||||
- Use `uv` for package management
|
||||
- Async throughout (asyncpg, async FastAPI endpoints)
|
||||
- Pydantic models for request/response validation
|
||||
- No py files at project root - maintain clean directory structure
|
||||
|
||||
### TypeScript (control-plane, clients)
|
||||
- Next.js with App Router for control plane
|
||||
- Tailwind CSS with shadcn/ui components
|
||||
|
||||
### Rust (CLI)
|
||||
- Async with tokio
|
||||
- reqwest for HTTP client
|
||||
- progenitor for API client generation
|
||||
|
||||
## Database
|
||||
|
||||
- PostgreSQL with pgvector extension
|
||||
- Schema managed via Alembic migrations in `hindsight-api/alembic/`, db migrations happen during api startup, no manual commands
|
||||
- Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
|
||||
253
NEW.md
253
NEW.md
|
|
@ -1,253 +0,0 @@
|
|||
<div align="center">
|
||||
|
||||
# Hindsight
|
||||
|
||||
**Agent Memory that Works Like Human Memory**
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/test.yml)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](https://pypi.org/project/hindsight-client/)
|
||||
[](https://www.npmjs.com/package/@vectorize-io/hindsight-client)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3klo21kua-VUCC_zHP5rIcXFB1_5yw6A)
|
||||
|
||||
[Documentation](https://vectorize-io.github.io/hindsight) • [Paper](./Hindsight.pdf) • [Examples](./examples)
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
## What is Hindsight?
|
||||
|
||||
Hindsight is an agent memory system built to create smarter agents that learn over time. It eliminates the shortcomings of alternative techniques such as RAG and knowledge graph.
|
||||
|
||||
Hindsight addresses common challenges that have frustrated AI engineers building agents to automate tasks and assist users with conversational interfaces. Many of these challenges stem directly from a lack of memory.
|
||||
|
||||
- **Inconsistency:** Agents complete tasks successfully one time, then fail when asked to complete the same task again. Memory gives the agent a mechanism to remember what worked and what didn't and to use that information to reduce errors and improve consistency.
|
||||
- **Hallucinations:** Long term memory can be seeded with external knowledge to ground agent behavior in reliable sources to augment training data.
|
||||
- **Cognitive Overload:** As workflows get complex, retrievals, tool calls, user messages and agent responses can grow to fill the context window leading to context rot. Short term memory optimization allows agents to reduce tokens and focus context by removing irrelevant details.
|
||||
|
||||
## How Hindsight Works
|
||||
|
||||
|
||||
|
||||
Hindsight organizes memory into four networks to mimic the way human memory works:
|
||||
|
||||
- **World:** Facts about the world ("The stove gets hot")
|
||||
- **Experiences:** Agent's own experiences ("I touched the stove and it really hurt")
|
||||
- **Opinion:** Beliefs with confidence scores ("I shouldn't touch the stove again" - .99 confidence)
|
||||
- **Observation:** Complex mental models derived by reflecting on facts and experiences ("Curling irons, ovens, and fire are also hot. I shouldn't touch those either.")
|
||||
|
||||
Hindsight provides three simple methods to interact with the system:
|
||||
|
||||
- **Retain:** Provide information to Hindsight that you want it to remember
|
||||
- **Recall:** Retrieve memories from Hindsight
|
||||
- **Reflect:** Reflect on memories and experiences to generate new observations and insights from existing memories.
|
||||
|
||||
Memories in Hindsight are stored in banks (e.g. memory banks). When memories are retained, they are transformed to construct a series of search indexes, time series data, and entity/relationship graphs.
|
||||
|
||||
---
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Docker (recommended)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=your-key
|
||||
|
||||
docker run -p 8888:8888 -p 9999:9999 \
|
||||
-e HINDSIGHT_API_LLM_PROVIDER=openai \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-e HINDSIGHT_API_LLM_MODEL=gpt-4o-mini \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight
|
||||
```
|
||||
|
||||
API: http://localhost:8888
|
||||
UI: http://localhost:9999
|
||||
|
||||
Install client:
|
||||
|
||||
```bash
|
||||
pip install hindsight-client
|
||||
# or
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
Python example:
|
||||
|
||||
```python
|
||||
from hindsight import HindsightClient
|
||||
|
||||
client = HindsightClient(base_url="http://localhost:8888")
|
||||
|
||||
# Store
|
||||
client.retain(bank_id="my-agent", content="Alice works at Google as a software engineer")
|
||||
|
||||
# Query
|
||||
results = client.recall(bank_id="my-agent", query="What does Alice do?")
|
||||
|
||||
# Reflect
|
||||
response = client.reflect(bank_id="my-agent", query="Tell me about Alice")
|
||||
print(response.text)
|
||||
```
|
||||
|
||||
### Python (embedded, no Docker)
|
||||
|
||||
```bash
|
||||
pip install hindsight-all
|
||||
```
|
||||
|
||||
```python
|
||||
import os
|
||||
from hindsight import HindsightServer, HindsightClient
|
||||
|
||||
with HindsightServer(
|
||||
llm_provider="openai",
|
||||
llm_model="gpt-4o-mini",
|
||||
llm_api_key=os.environ["OPENAI_API_KEY"]
|
||||
) as server:
|
||||
client = HindsightClient(base_url=server.url)
|
||||
client.retain(bank_id="my-agent", content="Alice works at Google")
|
||||
results = client.recall(bank_id="my-agent", query="Where does Alice work?")
|
||||
```
|
||||
|
||||
### TypeScript
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```typescript
|
||||
import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
await client.retain('my-agent', 'Alice loves hiking in Yosemite');
|
||||
const response = await client.recall('my-agent', 'What does Alice like?');
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Architecture & Operations
|
||||
|
||||
### Retain
|
||||
|
||||
The `retain` operation is used to push new memories into Hindsight. It tells Hindsight to _retain_ the information you pass in as an input.
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# Simple
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice works at Google as a software engineer"
|
||||
)
|
||||
|
||||
# With context and timestamp
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice got promoted to senior engineer",
|
||||
context="career update",
|
||||
timestamp="2025-06-15T10:00:00Z"
|
||||
)
|
||||
```
|
||||
|
||||
Behind the scenes, the retain operation uses an LLM to extract key facts, temporal data, entities, and relationships. It passes these through a normalization process to transform extracted data into canonical entities, time series, and search indexes along with metadata. These representations create the pathways for accurate memory retrieval in the recall and reflect operations.
|
||||
|
||||
|
||||
|
||||
### Recall
|
||||
|
||||
The recall operation is used to retrieve memories. These memories can come from any of the memory types (world, experiences, etc.)
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# Simple
|
||||
client.recall(bank_id="my-bank", query="What does Alice do?")
|
||||
|
||||
# Temporal
|
||||
results = client.recall(bank_id="my-bank", query="What happened in June?")
|
||||
|
||||
|
||||
```
|
||||
|
||||
Recall performs 4 retrieval strategies in parallel:
|
||||
- Semantic: Vector similarity
|
||||
- Keyword: BM25 exact matching
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||
|
||||
|
||||
The individual results from the retrievals are merged, then ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model.
|
||||
|
||||
The final output is trimmed as needed to fit within the token limit.
|
||||
|
||||
### Reflect
|
||||
|
||||
The reflect operation is used to perform a more thorough analysis of existing memories. This allows the agent to form new connections between memories which are then persisted as opinions and/or observations. When building agents, the reflect operation is a key capability to enable the agent to learn from its experiences.
|
||||
|
||||
For example, the `reflect` operation can be used to support use cases such as:
|
||||
|
||||
- An **AI Project Manager** reflecting on what risks need to be mitigated on a project.
|
||||
- A **Sales Agent** reflecting on why certain outreach messages have gotten responses while others haven't.
|
||||
- A **Support Agent** reflecting on opportunities where customers have questions not answered by current product documentation.
|
||||
|
||||
The `reflect` operation can also be used to handle on-demand question answering or analysis which require more deep thinking.
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
```
|
||||
|
||||
|
||||
|
||||
## Integrations
|
||||
|
||||
### Examples
|
||||
|
||||
[Examples Repo]([./examples](https://github.com/vectorize-io/hindsight-cookbook)) includes:
|
||||
|
||||
- Basic usage
|
||||
- Multi-session conversations
|
||||
- Temporal queries
|
||||
- Entity reasoning
|
||||
- Opinion tracking
|
||||
- Production setup (Docker Compose + monitoring)
|
||||
|
||||
---
|
||||
|
||||
## Resources
|
||||
|
||||
**Documentation:** [vectorize-io.github.io/hindsight](https://vectorize-io.github.io/hindsight)
|
||||
|
||||
**Clients:**
|
||||
- [Python](http://hindsight.vectorize.io/sdks/python)
|
||||
- [Node.js](http://hindsight.vectorize.io/sdks/nodejs)
|
||||
- [REST API](http://hindsight.vectorize.io/api-reference)
|
||||
|
||||
**Community:**
|
||||
- [Slack](https://join.slack.com/t/hindsight-space/shared_invite/zt-3klo21kua-VUCC_zHP5rIcXFB1_5yw6A)
|
||||
- [GitHub Issues](https://github.com/vectorize-io/hindsight/issues)
|
||||
|
||||
---
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING.md](./CONTRIBUTING.md).
|
||||
|
||||
## License
|
||||
|
||||
MIT — see [LICENSE](./LICENSE)
|
||||
|
||||
---
|
||||
|
||||
Built by [Vectorize.io](https://vectorize.io)
|
||||
|
|
@ -16,7 +16,8 @@ from .engine.search.trace import (
|
|||
SearchPhaseMetrics,
|
||||
)
|
||||
from .engine.search.tracer import SearchTracer
|
||||
from .engine.embeddings import Embeddings, SentenceTransformersEmbeddings
|
||||
from .engine.embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
from .engine.cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
|
||||
from .engine.llm_wrapper import LLMConfig
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -32,7 +33,11 @@ __all__ = [
|
|||
"SearchSummary",
|
||||
"SearchPhaseMetrics",
|
||||
"Embeddings",
|
||||
"SentenceTransformersEmbeddings",
|
||||
"LocalSTEmbeddings",
|
||||
"RemoteTEIEmbeddings",
|
||||
"CrossEncoderModel",
|
||||
"LocalSTCrossEncoder",
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.1.0"
|
||||
|
|
|
|||
|
|
@ -9,7 +9,8 @@ This package contains all the implementation details of the memory engine:
|
|||
|
||||
from .memory_engine import MemoryEngine
|
||||
from .db_utils import acquire_with_retry
|
||||
from .embeddings import Embeddings, SentenceTransformersEmbeddings
|
||||
from .embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
from .cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
|
||||
from .search.trace import (
|
||||
SearchTrace,
|
||||
QueryInfo,
|
||||
|
|
@ -29,7 +30,11 @@ __all__ = [
|
|||
"MemoryEngine",
|
||||
"acquire_with_retry",
|
||||
"Embeddings",
|
||||
"SentenceTransformersEmbeddings",
|
||||
"LocalSTEmbeddings",
|
||||
"RemoteTEIEmbeddings",
|
||||
"CrossEncoderModel",
|
||||
"LocalSTCrossEncoder",
|
||||
"RemoteTEICrossEncoder",
|
||||
"SearchTrace",
|
||||
"SearchTracer",
|
||||
"QueryInfo",
|
||||
|
|
|
|||
|
|
@ -2,13 +2,28 @@
|
|||
Cross-encoder abstraction for reranking.
|
||||
|
||||
Provides an interface for reranking with different backends.
|
||||
|
||||
Configuration via environment variables:
|
||||
- HINDSIGHT_API_RERANKER_PROVIDER: "local" (default) or "tei"
|
||||
|
||||
For local provider:
|
||||
- HINDSIGHT_API_RERANKER_LOCAL_MODEL: Model name (default: cross-encoder/ms-marco-MiniLM-L-6-v2)
|
||||
|
||||
For TEI provider:
|
||||
- HINDSIGHT_API_RERANKER_TEI_URL: TEI server URL (required)
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Tuple
|
||||
from typing import List, Tuple, Optional
|
||||
import logging
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default model for local cross-encoder
|
||||
DEFAULT_RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
|
|
@ -17,12 +32,18 @@ class CrossEncoderModel(ABC):
|
|||
Cross-encoders take query-document pairs and return relevance scores.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def load(self) -> None:
|
||||
"""
|
||||
Load the cross-encoder model.
|
||||
def provider_name(self) -> str:
|
||||
"""Return a human-readable name for this provider (e.g., 'local', 'tei')."""
|
||||
pass
|
||||
|
||||
This should be called during initialization to load the model
|
||||
@abstractmethod
|
||||
async def initialize(self) -> None:
|
||||
"""
|
||||
Initialize the cross-encoder model asynchronously.
|
||||
|
||||
This should be called during startup to load/connect to the model
|
||||
and avoid cold start latency on first predict() call.
|
||||
"""
|
||||
pass
|
||||
|
|
@ -41,11 +62,11 @@ class CrossEncoderModel(ABC):
|
|||
pass
|
||||
|
||||
|
||||
class SentenceTransformersCrossEncoder(CrossEncoderModel):
|
||||
class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Cross-encoder implementation using SentenceTransformers.
|
||||
Local cross-encoder implementation using SentenceTransformers.
|
||||
|
||||
Call load() during initialization to load the model and avoid cold starts.
|
||||
Call initialize() during startup to load the model and avoid cold starts.
|
||||
|
||||
Default model is cross-encoder/ms-marco-MiniLM-L-6-v2:
|
||||
- Fast inference (~80ms for 100 pairs on CPU)
|
||||
|
|
@ -53,18 +74,22 @@ class SentenceTransformersCrossEncoder(CrossEncoderModel):
|
|||
- Trained for passage re-ranking
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
|
||||
def __init__(self, model_name: Optional[str] = None):
|
||||
"""
|
||||
Initialize SentenceTransformers cross-encoder.
|
||||
Initialize local SentenceTransformers cross-encoder.
|
||||
|
||||
Args:
|
||||
model_name: Name of the CrossEncoder model to use.
|
||||
Default: cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.model_name = model_name or DEFAULT_RERANKER_MODEL
|
||||
self._model = None
|
||||
|
||||
def load(self) -> None:
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "local"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Load the cross-encoder model."""
|
||||
if self._model is not None:
|
||||
return
|
||||
|
|
@ -73,18 +98,18 @@ class SentenceTransformersCrossEncoder(CrossEncoderModel):
|
|||
from sentence_transformers import CrossEncoder
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"sentence-transformers is required for SentenceTransformersCrossEncoder. "
|
||||
"sentence-transformers is required for LocalSTCrossEncoder. "
|
||||
"Install it with: pip install sentence-transformers"
|
||||
)
|
||||
|
||||
logger.info(f"Loading cross-encoder model: {self.model_name}...")
|
||||
logger.info(f"Reranker: initializing local provider with model {self.model_name}")
|
||||
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
|
||||
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
|
||||
self._model = CrossEncoder(
|
||||
self.model_name,
|
||||
model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
|
||||
)
|
||||
logger.info("Cross-encoder model loaded")
|
||||
logger.info("Reranker: local provider initialized")
|
||||
|
||||
def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
|
||||
"""
|
||||
|
|
@ -97,6 +122,194 @@ class SentenceTransformersCrossEncoder(CrossEncoderModel):
|
|||
List of relevance scores (raw logits from the model)
|
||||
"""
|
||||
if self._model is None:
|
||||
self.load()
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
scores = self._model.predict(pairs, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, 'tolist') else list(scores)
|
||||
|
||||
|
||||
class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Remote cross-encoder implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
|
||||
|
||||
TEI supports reranking via the /rerank endpoint.
|
||||
See: https://github.com/huggingface/text-embeddings-inference
|
||||
|
||||
Note: The TEI server must be running a cross-encoder/reranker model.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
timeout: float = 30.0,
|
||||
batch_size: int = 32,
|
||||
max_retries: int = 3,
|
||||
retry_delay: float = 0.5,
|
||||
):
|
||||
"""
|
||||
Initialize remote TEI cross-encoder client.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of the TEI server (e.g., "http://localhost:8080")
|
||||
timeout: Request timeout in seconds (default: 30.0)
|
||||
batch_size: Maximum batch size for rerank requests (default: 32)
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5)
|
||||
"""
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.timeout = timeout
|
||||
self.batch_size = batch_size
|
||||
self.max_retries = max_retries
|
||||
self.retry_delay = retry_delay
|
||||
self._client: Optional[httpx.Client] = None
|
||||
self._model_id: Optional[str] = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "tei"
|
||||
|
||||
def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response:
|
||||
"""Make an HTTP request with automatic retries on transient errors."""
|
||||
import time
|
||||
last_error = None
|
||||
delay = self.retry_delay
|
||||
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
if method == "GET":
|
||||
response = self._client.get(url, **kwargs)
|
||||
else:
|
||||
response = self._client.post(url, **kwargs)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
|
||||
last_error = e
|
||||
if attempt < self.max_retries:
|
||||
logger.warning(f"TEI request failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
delay *= 2 # Exponential backoff
|
||||
except httpx.HTTPStatusError as e:
|
||||
# Retry on 5xx server errors
|
||||
if e.response.status_code >= 500 and attempt < self.max_retries:
|
||||
last_error = e
|
||||
logger.warning(f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
delay *= 2
|
||||
else:
|
||||
raise
|
||||
|
||||
raise last_error
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the HTTP client and verify server connectivity."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
logger.info(f"Reranker: initializing TEI provider at {self.base_url}")
|
||||
self._client = httpx.Client(timeout=self.timeout)
|
||||
|
||||
# Verify server is reachable and get model info
|
||||
try:
|
||||
response = self._request_with_retry("GET", f"{self.base_url}/info")
|
||||
info = response.json()
|
||||
self._model_id = info.get("model_id", "unknown")
|
||||
logger.info(f"Reranker: TEI provider initialized (model: {self._model_id})")
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}")
|
||||
|
||||
def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
|
||||
"""
|
||||
Score query-document pairs using the remote TEI reranker.
|
||||
|
||||
Args:
|
||||
pairs: List of (query, document) tuples to score
|
||||
|
||||
Returns:
|
||||
List of relevance scores
|
||||
"""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
all_scores = []
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(pairs), self.batch_size):
|
||||
batch = pairs[i:i + self.batch_size]
|
||||
|
||||
# TEI rerank endpoint expects query and texts separately
|
||||
# All pairs in a batch should have the same query for optimal performance
|
||||
# but we handle mixed queries by making separate requests per unique query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(batch):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
batch_scores = [0.0] * len(batch)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
try:
|
||||
response = self._request_with_retry(
|
||||
"POST",
|
||||
f"{self.base_url}/rerank",
|
||||
json={
|
||||
"query": query,
|
||||
"texts": texts,
|
||||
"return_text": False,
|
||||
},
|
||||
)
|
||||
results = response.json()
|
||||
|
||||
# TEI returns results sorted by score descending, with original index
|
||||
for result in results:
|
||||
original_idx = result["index"]
|
||||
score = result["score"]
|
||||
# Map back to batch position
|
||||
batch_scores[indices[original_idx]] = score
|
||||
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"TEI rerank request failed: {e}")
|
||||
|
||||
all_scores.extend(batch_scores)
|
||||
|
||||
return all_scores
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on environment variables.
|
||||
|
||||
Environment variables:
|
||||
- HINDSIGHT_API_RERANKER_PROVIDER: "local" (default) or "tei"
|
||||
|
||||
For local provider:
|
||||
- HINDSIGHT_API_RERANKER_LOCAL_MODEL: Model name (default: cross-encoder/ms-marco-MiniLM-L-6-v2)
|
||||
|
||||
For TEI provider:
|
||||
- HINDSIGHT_API_RERANKER_TEI_URL: TEI server URL (required)
|
||||
|
||||
Returns:
|
||||
Configured CrossEncoderModel instance
|
||||
"""
|
||||
provider = os.environ.get("HINDSIGHT_API_RERANKER_PROVIDER", "local").lower()
|
||||
|
||||
if provider == "tei":
|
||||
url = os.environ.get("HINDSIGHT_API_RERANKER_TEI_URL")
|
||||
if not url:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_RERANKER_TEI_URL is required when HINDSIGHT_API_RERANKER_PROVIDER is 'tei'"
|
||||
)
|
||||
return RemoteTEICrossEncoder(base_url=url)
|
||||
elif provider == "local":
|
||||
model = os.environ.get("HINDSIGHT_API_RERANKER_LOCAL_MODEL")
|
||||
model_name = model or DEFAULT_RERANKER_MODEL
|
||||
return LocalSTCrossEncoder(model_name=model_name)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei'"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -5,16 +5,31 @@ Provides an interface for generating embeddings with different backends.
|
|||
|
||||
IMPORTANT: All embeddings must produce 384-dimensional vectors to match
|
||||
the database schema (pgvector column defined as vector(384)).
|
||||
|
||||
Configuration via environment variables:
|
||||
- HINDSIGHT_API_EMBEDDINGS_PROVIDER: "local" (default) or "tei"
|
||||
|
||||
For local provider:
|
||||
- HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL: Model name (default: BAAI/bge-small-en-v1.5)
|
||||
|
||||
For TEI provider:
|
||||
- HINDSIGHT_API_EMBEDDINGS_TEI_URL: TEI server URL (required)
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List
|
||||
from typing import List, Optional
|
||||
import logging
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Fixed embedding dimension required by database schema
|
||||
EMBEDDING_DIMENSION = 384
|
||||
|
||||
# Default model for local embeddings
|
||||
DEFAULT_EMBEDDINGS_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""
|
||||
|
|
@ -24,12 +39,18 @@ class Embeddings(ABC):
|
|||
the database schema.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def load(self) -> None:
|
||||
"""
|
||||
Load the embedding model.
|
||||
def provider_name(self) -> str:
|
||||
"""Return a human-readable name for this provider (e.g., 'local', 'tei')."""
|
||||
pass
|
||||
|
||||
This should be called during initialization to load the model
|
||||
@abstractmethod
|
||||
async def initialize(self) -> None:
|
||||
"""
|
||||
Initialize the embedding model asynchronously.
|
||||
|
||||
This should be called during startup to load/connect to the model
|
||||
and avoid cold start latency on first encode() call.
|
||||
"""
|
||||
pass
|
||||
|
|
@ -48,29 +69,33 @@ class Embeddings(ABC):
|
|||
pass
|
||||
|
||||
|
||||
class SentenceTransformersEmbeddings(Embeddings):
|
||||
class LocalSTEmbeddings(Embeddings):
|
||||
"""
|
||||
Embeddings implementation using SentenceTransformers.
|
||||
Local embeddings implementation using SentenceTransformers.
|
||||
|
||||
Call load() during initialization to load the model and avoid cold starts.
|
||||
Call initialize() during startup to load the model and avoid cold starts.
|
||||
|
||||
Default model is BAAI/bge-small-en-v1.5 which produces 384-dimensional
|
||||
embeddings matching the database schema.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"):
|
||||
def __init__(self, model_name: Optional[str] = None):
|
||||
"""
|
||||
Initialize SentenceTransformers embeddings.
|
||||
Initialize local SentenceTransformers embeddings.
|
||||
|
||||
Args:
|
||||
model_name: Name of the SentenceTransformer model to use.
|
||||
Must produce 384-dimensional embeddings.
|
||||
Default: BAAI/bge-small-en-v1.5
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.model_name = model_name or DEFAULT_EMBEDDINGS_MODEL
|
||||
self._model = None
|
||||
|
||||
def load(self) -> None:
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "local"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Load the embedding model."""
|
||||
if self._model is not None:
|
||||
return
|
||||
|
|
@ -79,11 +104,11 @@ class SentenceTransformersEmbeddings(Embeddings):
|
|||
from sentence_transformers import SentenceTransformer
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"sentence-transformers is required for SentenceTransformersEmbeddings. "
|
||||
"sentence-transformers is required for LocalSTEmbeddings. "
|
||||
"Install it with: pip install sentence-transformers"
|
||||
)
|
||||
|
||||
logger.info(f"Loading embedding model: {self.model_name}...")
|
||||
logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
|
||||
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
|
||||
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
|
||||
self._model = SentenceTransformer(
|
||||
|
|
@ -100,7 +125,7 @@ class SentenceTransformersEmbeddings(Embeddings):
|
|||
f"Use a model that produces {EMBEDDING_DIMENSION}-dimensional embeddings."
|
||||
)
|
||||
|
||||
logger.info(f"Model loaded (embedding dim: {model_dim})")
|
||||
logger.info(f"Embeddings: local provider initialized (dim: {model_dim})")
|
||||
|
||||
def encode(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
|
|
@ -113,6 +138,166 @@ class SentenceTransformersEmbeddings(Embeddings):
|
|||
List of 384-dimensional embedding vectors
|
||||
"""
|
||||
if self._model is None:
|
||||
self.load()
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
|
||||
return [emb.tolist() for emb in embeddings]
|
||||
|
||||
|
||||
class RemoteTEIEmbeddings(Embeddings):
|
||||
"""
|
||||
Remote embeddings implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
|
||||
|
||||
TEI provides a high-performance inference server for embedding models.
|
||||
See: https://github.com/huggingface/text-embeddings-inference
|
||||
|
||||
The server should be running a model that produces 384-dimensional embeddings.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
timeout: float = 30.0,
|
||||
batch_size: int = 32,
|
||||
max_retries: int = 3,
|
||||
retry_delay: float = 0.5,
|
||||
):
|
||||
"""
|
||||
Initialize remote TEI embeddings client.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of the TEI server (e.g., "http://localhost:8080")
|
||||
timeout: Request timeout in seconds (default: 30.0)
|
||||
batch_size: Maximum batch size for embedding requests (default: 32)
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5)
|
||||
"""
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.timeout = timeout
|
||||
self.batch_size = batch_size
|
||||
self.max_retries = max_retries
|
||||
self.retry_delay = retry_delay
|
||||
self._client: Optional[httpx.Client] = None
|
||||
self._model_id: Optional[str] = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "tei"
|
||||
|
||||
def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response:
|
||||
"""Make an HTTP request with automatic retries on transient errors."""
|
||||
import time
|
||||
last_error = None
|
||||
delay = self.retry_delay
|
||||
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
if method == "GET":
|
||||
response = self._client.get(url, **kwargs)
|
||||
else:
|
||||
response = self._client.post(url, **kwargs)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
|
||||
last_error = e
|
||||
if attempt < self.max_retries:
|
||||
logger.warning(f"TEI request failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
delay *= 2 # Exponential backoff
|
||||
except httpx.HTTPStatusError as e:
|
||||
# Retry on 5xx server errors
|
||||
if e.response.status_code >= 500 and attempt < self.max_retries:
|
||||
last_error = e
|
||||
logger.warning(f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
delay *= 2
|
||||
else:
|
||||
raise
|
||||
|
||||
raise last_error
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the HTTP client and verify server connectivity."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
logger.info(f"Embeddings: initializing TEI provider at {self.base_url}")
|
||||
self._client = httpx.Client(timeout=self.timeout)
|
||||
|
||||
# Verify server is reachable and get model info
|
||||
try:
|
||||
response = self._request_with_retry("GET", f"{self.base_url}/info")
|
||||
info = response.json()
|
||||
self._model_id = info.get("model_id", "unknown")
|
||||
logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id})")
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}")
|
||||
|
||||
def encode(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Generate embeddings using the remote TEI server.
|
||||
|
||||
Args:
|
||||
texts: List of text strings to encode
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i:i + self.batch_size]
|
||||
|
||||
try:
|
||||
response = self._request_with_retry(
|
||||
"POST",
|
||||
f"{self.base_url}/embed",
|
||||
json={"inputs": batch},
|
||||
)
|
||||
batch_embeddings = response.json()
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"TEI embedding request failed: {e}")
|
||||
|
||||
return all_embeddings
|
||||
|
||||
|
||||
def create_embeddings_from_env() -> Embeddings:
|
||||
"""
|
||||
Create an Embeddings instance based on environment variables.
|
||||
|
||||
Environment variables:
|
||||
- HINDSIGHT_API_EMBEDDINGS_PROVIDER: "local" (default) or "tei"
|
||||
|
||||
For local provider:
|
||||
- HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL: Model name (default: BAAI/bge-small-en-v1.5)
|
||||
|
||||
For TEI provider:
|
||||
- HINDSIGHT_API_EMBEDDINGS_TEI_URL: TEI server URL (required)
|
||||
|
||||
Returns:
|
||||
Configured Embeddings instance
|
||||
"""
|
||||
provider = os.environ.get("HINDSIGHT_API_EMBEDDINGS_PROVIDER", "local").lower()
|
||||
|
||||
if provider == "tei":
|
||||
url = os.environ.get("HINDSIGHT_API_EMBEDDINGS_TEI_URL")
|
||||
if not url:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_TEI_URL is required when HINDSIGHT_API_EMBEDDINGS_PROVIDER is 'tei'"
|
||||
)
|
||||
return RemoteTEIEmbeddings(base_url=url)
|
||||
elif provider == "local":
|
||||
model = os.environ.get("HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL")
|
||||
model_name = model or DEFAULT_EMBEDDINGS_MODEL
|
||||
return LocalSTEmbeddings(model_name=model_name)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei'"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,8 +14,8 @@ from datetime import datetime, timedelta, timezone
|
|||
from typing import Any, Dict, List, Optional, Tuple, Union, TypedDict
|
||||
import asyncpg
|
||||
import asyncio
|
||||
from .embeddings import Embeddings, SentenceTransformersEmbeddings
|
||||
from .cross_encoder import CrossEncoderModel
|
||||
from .embeddings import Embeddings, create_embeddings_from_env
|
||||
from .cross_encoder import CrossEncoderModel, create_cross_encoder_from_env
|
||||
import time
|
||||
import numpy as np
|
||||
import uuid
|
||||
|
|
@ -124,7 +124,7 @@ class MemoryEngine:
|
|||
memory_llm_base_url: Base URL for the LLM API. Optional. Defaults based on provider:
|
||||
- groq: https://api.groq.com/openai/v1
|
||||
- ollama: http://localhost:11434/v1
|
||||
embeddings: Embeddings implementation to use. If not provided, uses SentenceTransformersEmbeddings
|
||||
embeddings: Embeddings implementation to use. If not provided, uses LocalSTEmbeddings
|
||||
cross_encoder: Cross-encoder model for reranking. If not provided, uses default when cross-encoder reranker is selected
|
||||
query_analyzer: Query analyzer implementation to use. If not provided, uses TransformerQueryAnalyzer
|
||||
pool_min_size: Minimum number of connections in the pool (default: 5)
|
||||
|
|
@ -184,11 +184,11 @@ class MemoryEngine:
|
|||
# Initialize entity resolver (will be created in initialize())
|
||||
self.entity_resolver = None
|
||||
|
||||
# Initialize embeddings
|
||||
# Initialize embeddings (from env vars if not provided)
|
||||
if embeddings is not None:
|
||||
self.embeddings = embeddings
|
||||
else:
|
||||
self.embeddings = SentenceTransformersEmbeddings("BAAI/bge-small-en-v1.5")
|
||||
self.embeddings = create_embeddings_from_env()
|
||||
|
||||
# Initialize query analyzer
|
||||
if query_analyzer is not None:
|
||||
|
|
@ -414,31 +414,40 @@ class MemoryEngine:
|
|||
if not was_already_running:
|
||||
self._pg0 = pg0
|
||||
|
||||
def load_embeddings():
|
||||
"""Load embedding model (CPU-bound)."""
|
||||
self.embeddings.load()
|
||||
async def init_embeddings():
|
||||
"""Initialize embedding model."""
|
||||
# For local providers, run in thread pool to avoid blocking event loop
|
||||
if self.embeddings.provider_name == "local":
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: asyncio.run(self.embeddings.initialize())
|
||||
)
|
||||
else:
|
||||
await self.embeddings.initialize()
|
||||
|
||||
def load_cross_encoder():
|
||||
"""Load cross-encoder model (CPU-bound)."""
|
||||
self._cross_encoder_reranker.cross_encoder.load()
|
||||
async def init_cross_encoder():
|
||||
"""Initialize cross-encoder model."""
|
||||
cross_encoder = self._cross_encoder_reranker.cross_encoder
|
||||
# For local providers, run in thread pool to avoid blocking event loop
|
||||
if cross_encoder.provider_name == "local":
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: asyncio.run(cross_encoder.initialize())
|
||||
)
|
||||
else:
|
||||
await cross_encoder.initialize()
|
||||
|
||||
def load_query_analyzer():
|
||||
"""Load query analyzer model (CPU-bound)."""
|
||||
self.query_analyzer.load()
|
||||
async def init_query_analyzer():
|
||||
"""Initialize query analyzer model."""
|
||||
# Query analyzer load is sync and CPU-bound
|
||||
await loop.run_in_executor(None, self.query_analyzer.load)
|
||||
|
||||
# Run pg0 and all model loads in parallel
|
||||
# pg0 is async (IO-bound), models are sync (CPU-bound in thread pool)
|
||||
# Use 3 workers to load all models concurrently
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
|
||||
# Start all tasks
|
||||
pg0_task = asyncio.create_task(start_pg0())
|
||||
embeddings_future = loop.run_in_executor(executor, load_embeddings)
|
||||
cross_encoder_future = loop.run_in_executor(executor, load_cross_encoder)
|
||||
query_analyzer_future = loop.run_in_executor(executor, load_query_analyzer)
|
||||
|
||||
# Wait for all to complete
|
||||
# Run pg0 and all model initializations in parallel
|
||||
await asyncio.gather(
|
||||
pg0_task, embeddings_future, cross_encoder_future, query_analyzer_future
|
||||
start_pg0(),
|
||||
init_embeddings(),
|
||||
init_cross_encoder(),
|
||||
init_query_analyzer(),
|
||||
)
|
||||
|
||||
# Run database migrations if enabled
|
||||
|
|
|
|||
|
|
@ -10,10 +10,8 @@ class CrossEncoderReranker:
|
|||
"""
|
||||
Neural reranking using a cross-encoder model.
|
||||
|
||||
Uses cross-encoder/ms-marco-MiniLM-L-6-v2 by default:
|
||||
- Fast inference (~80ms for 100 pairs on CPU)
|
||||
- Small model (80MB)
|
||||
- Trained for passage re-ranking
|
||||
Configured via environment variables (see cross_encoder.py).
|
||||
Default local model is cross-encoder/ms-marco-MiniLM-L-6-v2.
|
||||
"""
|
||||
|
||||
def __init__(self, cross_encoder=None):
|
||||
|
|
@ -21,14 +19,12 @@ class CrossEncoderReranker:
|
|||
Initialize cross-encoder reranker.
|
||||
|
||||
Args:
|
||||
cross_encoder: CrossEncoderReranker instance. If None, uses default
|
||||
SentenceTransformersCrossEncoder with ms-marco-MiniLM-L-6-v2
|
||||
(loaded lazily for faster startup)
|
||||
cross_encoder: CrossEncoderModel instance. If None, creates one from
|
||||
environment variables (defaults to local provider)
|
||||
"""
|
||||
if cross_encoder is None:
|
||||
from hindsight_api.engine.cross_encoder import SentenceTransformersCrossEncoder
|
||||
# Model is loaded lazily - call ensure_loaded() during initialize()
|
||||
cross_encoder = SentenceTransformersCrossEncoder()
|
||||
from hindsight_api.engine.cross_encoder import create_cross_encoder_from_env
|
||||
cross_encoder = create_cross_encoder_from_env()
|
||||
self.cross_encoder = cross_encoder
|
||||
|
||||
def rerank(
|
||||
|
|
|
|||
|
|
@ -1,46 +0,0 @@
|
|||
"""
|
||||
Debug script to test chunk extraction.
|
||||
"""
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from hindsight_api.engine.utils import extract_facts
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
import os
|
||||
|
||||
async def main():
|
||||
# Set up LLM config
|
||||
llm_config = LLMConfig.for_memory()
|
||||
|
||||
# Test content
|
||||
long_content = """
|
||||
Alice is a senior software engineer at TechCorp. She has been working there for 5 years.
|
||||
Alice specializes in distributed systems and has led the development of the company's
|
||||
microservices architecture. She is known for writing clean, well-documented code.
|
||||
|
||||
Bob joined the team last month as a junior developer. He is learning React and Node.js.
|
||||
Bob is enthusiastic and asks great questions during code reviews. He recently completed
|
||||
his first feature, which was a user authentication flow.
|
||||
|
||||
The team uses Kubernetes for container orchestration and deploys to AWS. They follow
|
||||
agile methodologies with two-week sprints. Code reviews are mandatory before merging.
|
||||
"""
|
||||
|
||||
# Extract facts and chunks
|
||||
facts, chunks = await extract_facts(
|
||||
text=long_content,
|
||||
event_date=datetime(2024, 1, 15),
|
||||
context="team overview",
|
||||
llm_config=llm_config
|
||||
)
|
||||
|
||||
print(f"\n=== Extracted {len(facts)} facts ===")
|
||||
for i, fact in enumerate(facts):
|
||||
print(f"{i+1}. {fact.fact[:100]}...")
|
||||
|
||||
print(f"\n=== Extracted {len(chunks)} chunks ===")
|
||||
for i, (chunk_text, fact_count) in enumerate(chunks):
|
||||
print(f"Chunk {i}: {fact_count} facts, {len(chunk_text)} chars")
|
||||
print(f" Text: {chunk_text[:100]}...")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
|
@ -1,69 +0,0 @@
|
|||
"""Test to verify mentioned_at uses event_date, not now()"""
|
||||
import asyncio
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
async def test_mentioned_at_uses_event_date():
|
||||
"""Verify that mentioned_at is set to event_date, not now()"""
|
||||
|
||||
# Use a date that's clearly not "now"
|
||||
past_date = datetime(2020, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
memory = MemoryEngine()
|
||||
await memory.initialize()
|
||||
|
||||
try:
|
||||
bank_id = "test_mentioned_at_debug"
|
||||
|
||||
# Store with explicit past event_date
|
||||
unit_ids = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alex went to the store.",
|
||||
context="test",
|
||||
event_date=past_date
|
||||
)
|
||||
|
||||
print(f"\n✅ Stored {len(unit_ids)} units")
|
||||
|
||||
# Recall and check mentioned_at
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="store",
|
||||
max_tokens=500
|
||||
)
|
||||
|
||||
print(f"✅ Found {len(result.results)} facts")
|
||||
|
||||
for i, fact in enumerate(result.results, 1):
|
||||
print(f"\nFact {i}:")
|
||||
print(f" Text: {fact.text[:80]}...")
|
||||
print(f" mentioned_at: {fact.mentioned_at}")
|
||||
print(f" occurred_start: {fact.occurred_start}")
|
||||
|
||||
# Parse mentioned_at
|
||||
if isinstance(fact.mentioned_at, str):
|
||||
mentioned_dt = datetime.fromisoformat(fact.mentioned_at.replace('Z', '+00:00'))
|
||||
else:
|
||||
mentioned_dt = fact.mentioned_at
|
||||
|
||||
# Check if mentioned_at matches our event_date
|
||||
time_diff = abs((mentioned_dt - past_date).total_seconds())
|
||||
|
||||
if time_diff < 60:
|
||||
print(f" ✅ mentioned_at correctly set to event_date")
|
||||
else:
|
||||
print(f" ❌ mentioned_at is {mentioned_dt}, expected {past_date}")
|
||||
print(f" Time difference: {time_diff} seconds")
|
||||
|
||||
# Check if it's close to now()
|
||||
now_diff = abs((mentioned_dt - datetime.now(timezone.utc)).total_seconds())
|
||||
if now_diff < 60:
|
||||
print(f" ⚠️ mentioned_at is using now() instead of event_date!")
|
||||
|
||||
await memory.delete_bank(bank_id)
|
||||
|
||||
finally:
|
||||
await memory.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(test_mentioned_at_uses_event_date())
|
||||
|
|
@ -8,9 +8,9 @@ import os
|
|||
import filelock
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
from hindsight_api import MemoryEngine, LLMConfig, SentenceTransformersEmbeddings
|
||||
from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings
|
||||
|
||||
from hindsight_api.engine.cross_encoder import SentenceTransformersCrossEncoder
|
||||
from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder
|
||||
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
|
||||
from hindsight_api.pg0 import EmbeddedPostgres
|
||||
|
||||
|
|
@ -111,14 +111,14 @@ def llm_config():
|
|||
@pytest.fixture(scope="session")
|
||||
def embeddings():
|
||||
|
||||
return SentenceTransformersEmbeddings("BAAI/bge-small-en-v1.5")
|
||||
return LocalSTEmbeddings()
|
||||
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def cross_encoder():
|
||||
|
||||
return SentenceTransformersCrossEncoder()
|
||||
return LocalSTCrossEncoder()
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def query_analyzer():
|
||||
|
|
|
|||
|
|
@ -80,19 +80,93 @@ export HINDSIGHT_API_MCP_ENABLED=true
|
|||
export HINDSIGHT_API_MCP_ENABLED=false
|
||||
```
|
||||
|
||||
### Embeddings Configuration
|
||||
|
||||
Configure the embeddings provider for semantic search. By default, uses local SentenceTransformers models.
|
||||
|
||||
| Variable | Description | Default | Required |
|
||||
|----------|-------------|---------|----------|
|
||||
| `HINDSIGHT_API_EMBEDDINGS_PROVIDER` | Provider: `local` or `tei` | `local` | No |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL` | Model name for local provider | `BAAI/bge-small-en-v1.5` | No |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_TEI_URL` | TEI server URL | - | Yes (if provider is `tei`) |
|
||||
|
||||
**Local Provider (Default)**
|
||||
|
||||
Uses SentenceTransformers to run embedding models locally. Good for development and smaller deployments.
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
export HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
```
|
||||
|
||||
**TEI Provider (HuggingFace Text Embeddings Inference)**
|
||||
|
||||
Uses a remote [TEI server](https://github.com/huggingface/text-embeddings-inference) for high-performance inference. Recommended for production deployments.
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=tei
|
||||
export HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
|
||||
```
|
||||
|
||||
:::warning
|
||||
All embedding models must produce 384-dimensional vectors to match the database schema.
|
||||
:::
|
||||
|
||||
### Reranker Configuration
|
||||
|
||||
Configure the cross-encoder reranker for improving search result relevance. By default, uses local SentenceTransformers models.
|
||||
|
||||
| Variable | Description | Default | Required |
|
||||
|----------|-------------|---------|----------|
|
||||
| `HINDSIGHT_API_RERANKER_PROVIDER` | Provider: `local` or `tei` | `local` | No |
|
||||
| `HINDSIGHT_API_RERANKER_LOCAL_MODEL` | Model name for local provider | `cross-encoder/ms-marco-MiniLM-L-6-v2` | No |
|
||||
| `HINDSIGHT_API_RERANKER_TEI_URL` | TEI server URL | - | Yes (if provider is `tei`) |
|
||||
|
||||
**Local Provider (Default)**
|
||||
|
||||
Uses SentenceTransformers CrossEncoder to run reranking locally.
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_RERANKER_PROVIDER=local
|
||||
export HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
```
|
||||
|
||||
**TEI Provider (HuggingFace Text Embeddings Inference)**
|
||||
|
||||
Uses a remote [TEI server](https://github.com/huggingface/text-embeddings-inference) with a reranker model.
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_RERANKER_PROVIDER=tei
|
||||
export HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
```
|
||||
|
||||
:::tip
|
||||
When using TEI, you can run separate servers for embeddings and reranking, or use a single server if it supports both operations with your chosen model.
|
||||
:::
|
||||
|
||||
## Configuration Files
|
||||
|
||||
### .env File
|
||||
|
||||
The Hindisight API will look for a `.env` file:
|
||||
The Hindsight API will look for a `.env` file:
|
||||
|
||||
```bash
|
||||
# .env
|
||||
|
||||
# Database
|
||||
HINDSIGHT_API_DATABASE_URL=postgresql://hindsight:hindsight_dev@localhost:5432/hindsight
|
||||
|
||||
# LLM
|
||||
HINDSIGHT_API_LLM_PROVIDER=groq
|
||||
HINDSIGHT_API_LLM_API_KEY=gsk_xxxxxxxxxxxx
|
||||
|
||||
# Embeddings (optional, defaults to local)
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
|
||||
# Reranker (optional, defaults to local)
|
||||
# HINDSIGHT_API_RERANKER_PROVIDER=local
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
```
|
||||
|
||||
---
|
||||
|
|
|
|||
Loading…
Reference in a new issue