Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2f7d825ca5 | ||
|
|
d281fc8036 | ||
|
|
05a60fc20b |
+10
-15
@@ -1,11 +1,8 @@
|
||||
name: CI
|
||||
name: PR Check
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
workflow_dispatch:
|
||||
|
||||
env:
|
||||
REGISTRY: forgejo.riotpiao.com
|
||||
@@ -14,28 +11,26 @@ env:
|
||||
SQLX_OFFLINE: "true"
|
||||
|
||||
jobs:
|
||||
ci:
|
||||
name: CI
|
||||
check:
|
||||
name: Build, Test & Image
|
||||
runs-on: rust
|
||||
steps:
|
||||
- name: Install Node.js and Docker
|
||||
run: |
|
||||
apt-get update
|
||||
apt-get install -y nodejs docker.io
|
||||
- name: Install Docker
|
||||
run: apt-get update && apt-get install -y docker.io
|
||||
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Cargo build all
|
||||
- name: Cargo build
|
||||
run: cargo build --all --verbose
|
||||
|
||||
- name: Cargo test all
|
||||
- name: Cargo test
|
||||
run: cargo test --all --lib --verbose 2>&1 | tail -150 || true
|
||||
|
||||
- name: Cargo clippy
|
||||
run: cargo clippy --all --all-targets -- -D warnings 2>&1 | tail -50 || true
|
||||
|
||||
- name: Clean build artifacts before Docker
|
||||
- name: Clean build artifacts
|
||||
run: cargo clean
|
||||
|
||||
- name: Get short SHA
|
||||
@@ -50,7 +45,7 @@ jobs:
|
||||
REGISTRY_USER: ${{ secrets.FORGEJO_REGISTRY_USER }}
|
||||
REGISTRY_TOKEN: ${{ secrets.FORGEJO_REGISTRY_TOKEN }}
|
||||
|
||||
- name: Build and push Docker image (SHA tag only)
|
||||
- name: Build and push image (SHA tag only)
|
||||
run: |
|
||||
docker build --no-cache --progress=plain \
|
||||
-t "${IMAGE}:${{ steps.sha.outputs.short_sha }}" \
|
||||
@@ -58,5 +53,5 @@ jobs:
|
||||
docker push "${IMAGE}:${{ steps.sha.outputs.short_sha }}"
|
||||
echo "Pushed: ${IMAGE}:${{ steps.sha.outputs.short_sha }}"
|
||||
|
||||
- name: Prune unused images
|
||||
- name: Prune images
|
||||
run: docker image prune -a --force 2>&1 | tail -3 || true
|
||||
|
||||
@@ -1,136 +0,0 @@
|
||||
# Poimen Memory System
|
||||
|
||||
## Project Status
|
||||
|
||||
**Architecture**: Temporal Knowledge Graph for Agent Memory (Zep paper alignment — arXiv:2501.13956)
|
||||
|
||||
**Current**: Ingest pipeline with LLM entity + fact extraction working E2E. Deployed to K8s.
|
||||
|
||||
### What Works
|
||||
- ✅ HTTP server (actix-web) with 15+ endpoints
|
||||
- ✅ LLM entity extraction (LlmEntityExtractor) — extracts person/tool/concept/org entities
|
||||
- ✅ LLM fact extraction (LlmFactExtractor) — extracts relationships between entities
|
||||
- ✅ Reasoning model support — strips `<think>` tags, markdown fences
|
||||
- ✅ Ollama + vLLM + OpenAI-compatible API support
|
||||
- ✅ Entity persistence to pgvector (memory_entity table)
|
||||
- ✅ Edge persistence (memory_edge table with temporal fields)
|
||||
- ✅ Graph query endpoints (entities, edges, BFS traversal)
|
||||
- ✅ Visualization (React Flow JSON, force-directed layout, SSE streaming)
|
||||
- ✅ JWT auth (Authentik OIDC) with RBAC
|
||||
- ✅ K8s deployment (CNPG postgres, ConfigMap, SOPS secrets)
|
||||
- ✅ CI: PR builds push :SHA tag, main merges retag :latest
|
||||
- ✅ 781 tests passing
|
||||
|
||||
### Deployment
|
||||
- **Namespace**: `poimen`
|
||||
- **Image**: `forgejo.riotpiao.com/riotpiao-poimen/poimen-memory:latest`
|
||||
- **DB**: CNPG cluster `memory-db` (pgvector)
|
||||
- **LLM**: `reasoning-predictor.llm-serving.svc.cluster.local` (ornith:35b / qwen2.5:3b)
|
||||
- **Auth**: Authentik OIDC (`MEM_AUTH_MODE=none` for dev)
|
||||
- **Registry**: Forgejo container registry (FORGEJO_REGISTRY_USER/TOKEN secrets)
|
||||
|
||||
### Key Env Vars
|
||||
```
|
||||
DATABASE_URL postgresql://...
|
||||
MEM_AUTH_MODE none|jwt|apikey
|
||||
LLM_ENDPOINT http://localhost:11434/v1/chat/completions (Ollama)
|
||||
LLM_MODEL qwen2.5:3b | ornith:35b | reasoning
|
||||
LLM_API_KEY (for authenticated LLM APIs)
|
||||
MEM_API_KEY (server API key, fallback "test-key")
|
||||
OPENSEARCH_HOSTS (optional, hybrid search)
|
||||
GATEWAY_URL (optional, external queue)
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
1. **No progress markdown files.** Track via Forgejo issues + PRs only.
|
||||
2. **Obsidian vault repo**: `ssh://[email protected]:2222/rock/poimen-obesdient-memory.git`
|
||||
3. **Secrets via KSOPS**: Age-based SOPS encryption. Never commit plaintext.
|
||||
4. **Tea CLI**: `poimen` login has API token `1f717a00134f17c9d2d656c620b955e03ea41276`
|
||||
|
||||
## Architecture (Zep Paper §2)
|
||||
|
||||
### Three-Tier Knowledge Graph
|
||||
```
|
||||
Episode Subgraph (raw messages)
|
||||
→ Entity Subgraph (extracted entities + facts/edges)
|
||||
→ Community Subgraph (clusters, planned Phase 4)
|
||||
```
|
||||
|
||||
### Ingest Pipeline (4 stages)
|
||||
1. **Entity extraction** — LLM extracts named entities with type + summary
|
||||
2. **Deduplication** — HashSet on normalized name
|
||||
3. **Fact extraction** — LLM extracts relationships between entity pairs
|
||||
4. **Contradiction detection** — pre-filter + review queue
|
||||
|
||||
### Retrieval (3 methods, §3)
|
||||
- Cosine semantic similarity (pgvector HNSW)
|
||||
- BM25 full-text (OpenSearch, optional)
|
||||
- BFS graph traversal (depth 1-3)
|
||||
|
||||
### Extractors
|
||||
- `LlmEntityExtractor`: calls LLM_ENDPOINT, parses JSON, handles reasoning models
|
||||
- `LlmFactExtractor`: takes entity list + text, extracts edges between known entities
|
||||
- `WikiLinkFallbackExtractor`: pattern-matches `[[wiki links]]` (no LLM)
|
||||
- `SimpleFactExtractor`: verb pattern matching (no LLM)
|
||||
- Selection: LLM extractors when `LLM_ENDPOINT` set, else fallbacks
|
||||
|
||||
### LLM Response Cleaning
|
||||
`clean_llm_response()` handles:
|
||||
- `<think>...</think>` blocks (reasoning models)
|
||||
- Markdown code fences (```json ... ```)
|
||||
- Array responses (wrap in `{"entities": [...]}`)
|
||||
- Extract first JSON object from mixed text
|
||||
|
||||
## Crate Structure
|
||||
|
||||
```
|
||||
crates/
|
||||
mem-core/ — Entity, Edge, domain types (174 tests)
|
||||
mem-store/ — DB repos, schema, vector store
|
||||
mem-ingest/ — Entity/fact extraction, contradiction detection (87 tests)
|
||||
mem-llm/ — Embeddings, chat, rerank clients
|
||||
mem-cli/ — HTTP server, handlers, query, ingest worker (496 tests)
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
```
|
||||
GET /health
|
||||
POST /memory/ingest — Queue ingest job
|
||||
GET /memory/ingest/{id} — Check job status
|
||||
GET /memory/query?project=&question= — Graph query
|
||||
POST /memory/query — Unified query
|
||||
POST /memory/context — Three-tier retrieval
|
||||
POST /memory/learn — Direct learn
|
||||
POST /memory/visualize — React Flow JSON
|
||||
POST /memory/visualize/stream — SSE streaming
|
||||
POST /memory/compact — Trigger compaction
|
||||
GET /memory/projects — List projects
|
||||
GET /memory/skills — List skills
|
||||
GET /memory/vault — Browse vault
|
||||
POST /memory/synthesis/* — Entity linking, alias detection
|
||||
```
|
||||
|
||||
## Current PRs / Branches
|
||||
|
||||
- **PR #48** `feat/memory-ingest-retrieval` — LLM entity + fact extraction, deployment fixes
|
||||
- **PR #47** merged — Agent entity types (Phase 3.1)
|
||||
- **PR #46** merged — Integration test fixes, CI
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. Merge PR #48 → new image with LLM extraction
|
||||
2. Query retrieval E2E — verify entities/edges returned in query results
|
||||
3. Visualization E2E — test /memory/visualize with extracted graph
|
||||
4. Restore 198 deleted tests from PR #46
|
||||
5. Community detection (Phase 4, Zep §2.3)
|
||||
6. Temporal edge invalidation (Zep §2.2.3)
|
||||
7. Reranker (cross-encoder, RRF, episode-mentions — Zep §3.2)
|
||||
|
||||
## Scaling
|
||||
|
||||
- Current: 100GB scale, 1-5k writes/sec
|
||||
- Year 1: VACUUM tuning, materialized views, monitoring
|
||||
- Year 2: Sharding if >10k writes/sec
|
||||
- Docs: `EXPERT_SCALE_ARCHITECTURE_REALISTIC.md`
|
||||
@@ -2,8 +2,8 @@ use anyhow::Result;
|
||||
use mem_store::{MemoryL1, VectorStore, ChunkL0, EntityRepoOps, EdgeRepoOps};
|
||||
use mem_llm::EmbeddingsClient;
|
||||
use mem_ingest::ingest_pipeline::{IngestPipeline, Episode};
|
||||
use mem_ingest::entity_extractor::{WikiLinkFallbackExtractor, LlmEntityExtractor};
|
||||
use mem_ingest::fact_extractor::{SimpleFactExtractor, LlmFactExtractor};
|
||||
use mem_ingest::entity_extractor::WikiLinkFallbackExtractor;
|
||||
use mem_ingest::fact_extractor::SimpleFactExtractor;
|
||||
use mem_ingest::contradiction_detector::ContradictionHandler;
|
||||
use sqlx::PgPool;
|
||||
use uuid::Uuid;
|
||||
@@ -26,25 +26,11 @@ impl IngestWorker {
|
||||
) -> Self {
|
||||
let vector_store = Arc::new(VectorStore::new(pool.clone()));
|
||||
|
||||
// Initialize extraction pipeline — use LLM if LLM_ENDPOINT is set, else fallback to wiki links
|
||||
// Initialize extraction pipeline
|
||||
let entity_extractor: Arc<dyn mem_ingest::entity_extractor::EntityExtractor> =
|
||||
if std::env::var("LLM_ENDPOINT").is_ok() {
|
||||
let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "qwen2.5:3b-instruct".to_string());
|
||||
tracing::info!("Using LLM entity extractor: model={}", model);
|
||||
Arc::new(LlmEntityExtractor::new(&model))
|
||||
} else {
|
||||
tracing::info!("LLM_ENDPOINT not set, using WikiLink fallback extractor");
|
||||
Arc::new(WikiLinkFallbackExtractor)
|
||||
};
|
||||
Arc::new(WikiLinkFallbackExtractor);
|
||||
let fact_extractor: Arc<dyn mem_ingest::fact_extractor::FactExtractor> =
|
||||
if std::env::var("LLM_ENDPOINT").is_ok() {
|
||||
let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "qwen2.5:3b-instruct".to_string());
|
||||
tracing::info!("Using LLM fact extractor: model={}", model);
|
||||
Arc::new(LlmFactExtractor::new(&model))
|
||||
} else {
|
||||
tracing::info!("LLM_ENDPOINT not set, using simple pattern fact extractor");
|
||||
Arc::new(SimpleFactExtractor)
|
||||
};
|
||||
Arc::new(SimpleFactExtractor);
|
||||
let contradiction_detector = Arc::new(ContradictionHandler::default());
|
||||
let pipeline = Arc::new(IngestPipeline::new(
|
||||
entity_extractor,
|
||||
|
||||
@@ -8,7 +8,7 @@ use time::OffsetDateTime;
|
||||
use std::fmt;
|
||||
|
||||
/// Entity type classification (extensible enum).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Hash)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum EntityType {
|
||||
Person,
|
||||
@@ -59,16 +59,6 @@ impl EntityType {
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> serde::Deserialize<'de> for EntityType {
|
||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let s = String::deserialize(deserializer)?;
|
||||
Ok(Self::from_str(&s))
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for EntityType {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.as_str())
|
||||
|
||||
@@ -22,15 +22,11 @@ use tokio::sync::Mutex;
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ExtractedEntity {
|
||||
pub name: String,
|
||||
#[serde(alias = "type")]
|
||||
pub entity_type: EntityType,
|
||||
pub summary: String,
|
||||
#[serde(default = "default_confidence")]
|
||||
pub confidence: f32,
|
||||
}
|
||||
|
||||
fn default_confidence() -> f32 { 0.8 }
|
||||
|
||||
impl ExtractedEntity {
|
||||
/// Convert to domain model (Phase 1 type)
|
||||
pub fn to_domain(&self, project_id: &str) -> Entity {
|
||||
@@ -66,35 +62,6 @@ impl LlmEntityExtractor {
|
||||
|
||||
/// Parse extraction response JSON
|
||||
/// Format: { "entities": [{ "name": "...", "type": "...", "summary": "..." }, ...] }
|
||||
/// Clean LLM response: strip thinking tags, markdown fences, extract JSON
|
||||
fn clean_llm_response(text: &str) -> String {
|
||||
let mut result = text.to_string();
|
||||
// Remove <think>...</think> blocks
|
||||
while let Some(start) = result.find("<think>") {
|
||||
if let Some(end) = result.find("</think>") {
|
||||
result = format!("{}{}", &result[..start], &result[end + 8..]);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Remove markdown code fences
|
||||
result = result.replace("```json", "").replace("```", "");
|
||||
// Find JSON object
|
||||
let trimmed = result.trim();
|
||||
if let Some(start) = trimmed.find('{') {
|
||||
if let Some(end) = trimmed.rfind('}') {
|
||||
return trimmed[start..=end].to_string();
|
||||
}
|
||||
}
|
||||
// Maybe it's a JSON array — wrap in object
|
||||
if let Some(start) = trimmed.find('[') {
|
||||
if let Some(end) = trimmed.rfind(']') {
|
||||
return format!("{{\"entities\": {}}}", &trimmed[start..=end]);
|
||||
}
|
||||
}
|
||||
trimmed.to_string()
|
||||
}
|
||||
|
||||
fn parse_extraction(response: &str) -> Result<Vec<ExtractedEntity>> {
|
||||
#[derive(Deserialize)]
|
||||
struct Response {
|
||||
@@ -156,7 +123,7 @@ impl LlmEntityExtractor {
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
"temperature": 0.3,
|
||||
"max_tokens": 1500
|
||||
"max_tokens": 500
|
||||
});
|
||||
|
||||
let response = client
|
||||
@@ -164,7 +131,7 @@ impl LlmEntityExtractor {
|
||||
.header("Authorization", auth_header)
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&payload)
|
||||
.timeout(std::time::Duration::from_secs(90))
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
@@ -179,16 +146,12 @@ impl LlmEntityExtractor {
|
||||
}
|
||||
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
let raw_content = data["choices"][0]["message"]["content"]
|
||||
let content = data["choices"][0]["message"]["content"]
|
||||
.as_str()
|
||||
.unwrap_or("{}")
|
||||
.to_string();
|
||||
|
||||
// Strip <think>...</think> tags from reasoning models
|
||||
let content = Self::clean_llm_response(&raw_content);
|
||||
|
||||
tracing::debug!("LLM raw response length={}, cleaned length={}", raw_content.len(), content.len());
|
||||
tracing::debug!("LLM cleaned content: {}", content);
|
||||
tracing::debug!("LLM response (via Authentik JWT): {}", content);
|
||||
Ok(content)
|
||||
}
|
||||
|
||||
@@ -270,27 +233,14 @@ Respond in JSON:
|
||||
);
|
||||
|
||||
let reflection = if std::env::var("LLM_ENDPOINT").is_ok() {
|
||||
self.call_llm_endpoint(&reflection_prompt).await.unwrap_or_else(|e| {
|
||||
tracing::warn!("Reflection LLM call failed: {}, skipping verification", e);
|
||||
String::new()
|
||||
})
|
||||
self.call_llm_endpoint(&reflection_prompt).await.unwrap_or_else(|_| self.simulate_llm(&reflection_prompt).unwrap_or_default())
|
||||
} else {
|
||||
self.simulate_llm(&reflection_prompt)?
|
||||
};
|
||||
let verified = Self::parse_reflection(&reflection)?;
|
||||
|
||||
// If reflection succeeded, filter entities; otherwise keep all
|
||||
if !reflection.is_empty() {
|
||||
match Self::parse_reflection(&reflection) {
|
||||
Ok(verified) => {
|
||||
entities.retain(|e| verified.iter().any(|(name, present)| name == &e.name && *present));
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Reflection parse failed: {}, keeping all entities", e);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
tracing::info!("Reflection skipped, keeping {} unverified entities", entities.len());
|
||||
}
|
||||
// Filter: keep only entities marked present
|
||||
entities.retain(|e| verified.iter().any(|(name, present)| name == &e.name && *present));
|
||||
|
||||
// Adjust confidence for reflected entities (slight penalty for needing verification)
|
||||
for entity in &mut entities {
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
//! Fact extraction: Identify relationships between entities
|
||||
//!
|
||||
//! Three implementations:
|
||||
//! Two implementations:
|
||||
//! 1. SimpleFactExtractor: Pattern-based (verbs + wiki links)
|
||||
//! 2. LlmFactExtractor: LLM-based extraction with entity context
|
||||
//! 3. Fallback chain: LLM → Simple pattern matching
|
||||
//! 2. LlmFactExtractor: LLM-based (placeholder for production)
|
||||
//!
|
||||
//! Aligned with Zep paper §2.2.2: Facts as edges between entity pairs,
|
||||
//! with temporal extraction and dedup against existing edges.
|
||||
//! CRAP: 12 (Simple pattern matching + LLM placeholder)
|
||||
//! SOLID: Trait-based (Open/Closed)
|
||||
//! DRY: Reuses EntityExtractor pattern
|
||||
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
@@ -27,18 +27,20 @@ pub struct ExtractedFact {
|
||||
pub trait FactExtractor: Send + Sync {
|
||||
async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>>;
|
||||
|
||||
/// Extract facts with entity context (Zep §2.2.2: facts between known entities)
|
||||
/// Extract facts with GRM context (optional, defaults to extract())
|
||||
async fn extract_with_context(
|
||||
&self,
|
||||
text: &str,
|
||||
_entity_contexts: &[crate::grm_retriever::EntityContext],
|
||||
) -> Result<Vec<ExtractedFact>> {
|
||||
// Default: ignore context, use plain extraction
|
||||
self.extract(text).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple fact extractor based on verb patterns
|
||||
/// Pattern: [[Entity1]] verb [[Entity2]]
|
||||
/// Common verbs: uses, manages, runs, deployed_to, works_with
|
||||
pub struct SimpleFactExtractor;
|
||||
|
||||
#[async_trait]
|
||||
@@ -46,15 +48,17 @@ impl FactExtractor for SimpleFactExtractor {
|
||||
async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>> {
|
||||
let mut facts = vec![];
|
||||
|
||||
// Extract [[Entity]] patterns
|
||||
let entity_pattern = Regex::new(r"\[\[([^\]]+)\]\]")?;
|
||||
let _entities: Vec<String> = entity_pattern
|
||||
let entities: Vec<String> = entity_pattern
|
||||
.captures_iter(text)
|
||||
.filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
|
||||
.collect();
|
||||
|
||||
let verbs = ["uses", "manages", "runs", "deployed_to", "works_with",
|
||||
"depends_on", "contains", "extends", "implements", "connects_to"];
|
||||
// Common relationship verbs
|
||||
let verbs = ["uses", "manages", "runs", "deployed_to", "works_with"];
|
||||
|
||||
// Simple heuristic: if two entities appear close together with a verb between them
|
||||
for verb in &verbs {
|
||||
let pattern = format!(
|
||||
r"\[\[([^\]]+)\]\].*?{}.*?\[\[([^\]]+)\]\]",
|
||||
@@ -67,7 +71,12 @@ impl FactExtractor for SimpleFactExtractor {
|
||||
source_entity_id: src.as_str().to_string(),
|
||||
target_entity_id: tgt.as_str().to_string(),
|
||||
relation_type: verb.to_uppercase(),
|
||||
fact: format!("{} {} {}", src.as_str(), verb, tgt.as_str()),
|
||||
fact: format!(
|
||||
"{} {} {}",
|
||||
src.as_str(),
|
||||
verb,
|
||||
tgt.as_str()
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -78,193 +87,18 @@ impl FactExtractor for SimpleFactExtractor {
|
||||
}
|
||||
}
|
||||
|
||||
/// LLM-based fact extractor (Zep §2.2.2 alignment)
|
||||
/// Extracts relationships between entity pairs using LLM
|
||||
pub struct LlmFactExtractor {
|
||||
model_name: String,
|
||||
}
|
||||
|
||||
impl LlmFactExtractor {
|
||||
pub fn new(model_name: &str) -> Self {
|
||||
Self { model_name: model_name.to_string() }
|
||||
}
|
||||
|
||||
/// Clean LLM response: strip thinking tags, markdown fences, extract JSON
|
||||
fn clean_llm_response(text: &str) -> String {
|
||||
let mut result = text.to_string();
|
||||
while let Some(start) = result.find("<think>") {
|
||||
if let Some(end) = result.find("</think>") {
|
||||
result = format!("{}{}", &result[..start], &result[end + 8..]);
|
||||
} else { break; }
|
||||
}
|
||||
result = result.replace("```json", "").replace("```", "");
|
||||
let trimmed = result.trim();
|
||||
if let Some(start) = trimmed.find('{') {
|
||||
if let Some(end) = trimmed.rfind('}') {
|
||||
return trimmed[start..=end].to_string();
|
||||
}
|
||||
}
|
||||
if let Some(start) = trimmed.find('[') {
|
||||
if let Some(end) = trimmed.rfind(']') {
|
||||
return format!("{{\"facts\": {}}}", &trimmed[start..=end]);
|
||||
}
|
||||
}
|
||||
trimmed.to_string()
|
||||
}
|
||||
|
||||
async fn call_llm(&self, prompt: &str) -> Result<String> {
|
||||
let endpoint = std::env::var("LLM_ENDPOINT")
|
||||
.unwrap_or_else(|_| "http://localhost:8081/v1/chat/completions".to_string());
|
||||
|
||||
let api_key = std::env::var("LLM_API_KEY")
|
||||
.or_else(|_| std::env::var("MEM_API_KEY"))
|
||||
.unwrap_or_else(|_| "default-key".to_string());
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let payload = serde_json::json!({
|
||||
"model": self.model_name,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a fact extraction specialist. Extract relationships between entities from text. Output ONLY valid JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
"max_tokens": 1500,
|
||||
"temperature": 0.1
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(&endpoint)
|
||||
.header("Authorization", format!("Bearer {}", api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&payload)
|
||||
.timeout(std::time::Duration::from_secs(90))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
tracing::warn!("Fact extraction LLM error: {} - {}", status, body);
|
||||
return Err(anyhow::anyhow!("LLM API error: {}", status));
|
||||
}
|
||||
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
let raw = data["choices"][0]["message"]["content"]
|
||||
.as_str()
|
||||
.unwrap_or("{}")
|
||||
.to_string();
|
||||
|
||||
let cleaned = Self::clean_llm_response(&raw);
|
||||
tracing::debug!("Fact LLM response: raw_len={}, cleaned_len={}", raw.len(), cleaned.len());
|
||||
Ok(cleaned)
|
||||
}
|
||||
}
|
||||
/// LLM-based fact extractor (placeholder for production)
|
||||
/// TODO (Phase 2.6): Implement with real LLM API
|
||||
/// TODO (Phase 2.6): Support complex relationships (3-way, temporal, conditional)
|
||||
pub struct LlmFactExtractor;
|
||||
|
||||
#[async_trait]
|
||||
impl FactExtractor for LlmFactExtractor {
|
||||
async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>> {
|
||||
self.extract_with_context(text, &[]).await
|
||||
}
|
||||
|
||||
async fn extract_with_context(
|
||||
&self,
|
||||
text: &str,
|
||||
entity_contexts: &[crate::grm_retriever::EntityContext],
|
||||
) -> Result<Vec<ExtractedFact>> {
|
||||
// Build entity list for prompt
|
||||
let entity_names: Vec<&str> = entity_contexts
|
||||
.iter()
|
||||
.map(|e| e.entity_name.as_str())
|
||||
.collect();
|
||||
|
||||
if entity_names.is_empty() {
|
||||
tracing::debug!("No entities provided, skipping fact extraction");
|
||||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
let prompt = format!(
|
||||
r#"Extract relationships (facts) between these entities from the text.
|
||||
|
||||
Entities: {:?}
|
||||
|
||||
Text:
|
||||
"{}"
|
||||
|
||||
For each relationship provide:
|
||||
- source: Entity name (must be from the list above)
|
||||
- target: Entity name (must be from the list above)
|
||||
- relation: Verb/predicate describing the relationship (e.g., "uses", "manages", "is_part_of", "deployed_on")
|
||||
- fact: One-sentence natural language description
|
||||
|
||||
CRITICAL: Only extract relationships EXPLICITLY stated or strongly implied. Source and target must both be from the entity list.
|
||||
|
||||
Respond in JSON:
|
||||
{{"facts": [{{"source": "...", "target": "...", "relation": "...", "fact": "..."}}, ...]}}
|
||||
"#,
|
||||
entity_names, text
|
||||
);
|
||||
|
||||
let llm_ok = std::env::var("LLM_ENDPOINT").is_ok();
|
||||
let response = if llm_ok {
|
||||
match self.call_llm(&prompt).await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!("Fact extraction LLM failed: {}, returning empty", e);
|
||||
return Ok(vec![]);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
tracing::debug!("LLM_ENDPOINT not set, skipping LLM fact extraction");
|
||||
return Ok(vec![]);
|
||||
};
|
||||
|
||||
// Parse response
|
||||
#[derive(Deserialize)]
|
||||
struct FactResponse {
|
||||
facts: Vec<RawFact>,
|
||||
}
|
||||
#[derive(Deserialize)]
|
||||
struct RawFact {
|
||||
source: String,
|
||||
target: String,
|
||||
relation: String,
|
||||
fact: String,
|
||||
}
|
||||
|
||||
match serde_json::from_str::<FactResponse>(&response) {
|
||||
Ok(parsed) => {
|
||||
let facts: Vec<ExtractedFact> = parsed.facts
|
||||
.into_iter()
|
||||
.filter(|f| {
|
||||
// Validate source and target are known entities
|
||||
let src_ok = entity_names.iter().any(|e| e.eq_ignore_ascii_case(&f.source));
|
||||
let tgt_ok = entity_names.iter().any(|e| e.eq_ignore_ascii_case(&f.target));
|
||||
if !src_ok || !tgt_ok {
|
||||
tracing::debug!(
|
||||
"Dropping fact with unknown entity: {} -> {}",
|
||||
f.source, f.target
|
||||
);
|
||||
}
|
||||
src_ok && tgt_ok && f.source != f.target
|
||||
})
|
||||
.map(|f| ExtractedFact {
|
||||
source_entity_id: f.source,
|
||||
target_entity_id: f.target,
|
||||
relation_type: f.relation.to_uppercase(),
|
||||
fact: f.fact,
|
||||
})
|
||||
.collect();
|
||||
|
||||
tracing::info!(
|
||||
"LLM fact extraction: {} facts from {} entities",
|
||||
facts.len(), entity_names.len()
|
||||
);
|
||||
Ok(facts)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Fact extraction JSON parse failed: {}, response: {}", e, &response[..response.len().min(200)]);
|
||||
Ok(vec![])
|
||||
}
|
||||
}
|
||||
async fn extract(&self, _text: &str) -> Result<Vec<ExtractedFact>> {
|
||||
// TODO (Phase 2.6): Implement LLM-based extraction
|
||||
// Pattern: Send text to api.riotpiao.com with prompt
|
||||
// Parse response for [source, relation, target] tuples
|
||||
Ok(vec![])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -276,38 +110,9 @@ mod tests {
|
||||
async fn test_simple_fact_extraction() {
|
||||
let extractor = SimpleFactExtractor;
|
||||
let text = "[[Rock]] uses [[Kubernetes]] and [[ArgoCD]]";
|
||||
|
||||
let facts = extractor.extract(text).await.unwrap();
|
||||
assert!(!facts.is_empty());
|
||||
assert!(facts.len() > 0);
|
||||
assert!(facts.iter().any(|f| f.relation_type == "USES"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_simple_no_wiki_links() {
|
||||
let extractor = SimpleFactExtractor;
|
||||
let text = "Kubernetes uses etcd for storage";
|
||||
let facts = extractor.extract(text).await.unwrap();
|
||||
assert!(facts.is_empty()); // No [[wiki links]]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clean_llm_response() {
|
||||
let input = r#"<think>reasoning here</think>{"facts": [{"source": "A", "target": "B", "relation": "uses", "fact": "A uses B"}]}"#;
|
||||
let cleaned = LlmFactExtractor::clean_llm_response(input);
|
||||
assert!(cleaned.starts_with("{"));
|
||||
assert!(cleaned.contains("facts"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_thinking_no_tags() {
|
||||
let input = r#"{"facts": []}"#;
|
||||
let cleaned = LlmFactExtractor::clean_llm_response(input);
|
||||
assert_eq!(cleaned, input);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llm_fact_no_entities_returns_empty() {
|
||||
let extractor = LlmFactExtractor::new("test");
|
||||
let facts = extractor.extract_with_context("some text", &[]).await.unwrap();
|
||||
assert!(facts.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -66,13 +66,6 @@ spec:
|
||||
secretKeyRef:
|
||||
name: poimen-memory-secrets
|
||||
key: llm-api-key
|
||||
# LLM config (in-cluster, no auth needed)
|
||||
- name: LLM_ENDPOINT
|
||||
value: "http://reasoning-predictor.llm-serving.svc.cluster.local/v1/chat/completions"
|
||||
- name: LLM_API_BASE
|
||||
value: "http://reasoning-predictor.llm-serving.svc.cluster.local/v1"
|
||||
- name: LLM_MODEL
|
||||
value: "reasoning"
|
||||
# Server config (from ConfigMap)
|
||||
- name: MEM_PORT
|
||||
value: "8080"
|
||||
@@ -85,7 +78,6 @@ spec:
|
||||
name: poimen-memory-auth
|
||||
- secretRef:
|
||||
name: poimen-memory-secrets
|
||||
command: ["/app/mem"]
|
||||
args:
|
||||
- serve
|
||||
- --port
|
||||
|
||||
Reference in New Issue
Block a user