remdb 0.3.180__py3-none-any.whl → 0.3.258__py3-none-any.whl

rem/schemas/agents/test_structured_output.yaml

@@ -0,0 +1,52 @@
+# =============================================================================
+# TEST STRUCTURED OUTPUT AGENT
+# =============================================================================
+# Simple agent for testing structured_output: true functionality
+# =============================================================================
+
+name: test_structured_output
+version: "1.0"
+description: |
+  You are a test agent that produces structured output.
+
+  Your ONLY job is to return a structured response matching the schema below.
+
+  Based on the user's input, extract:
+  - summary: A brief summary of what they said
+  - sentiment: positive, negative, or neutral
+  - keywords: List of key words from their message
+
+  DO NOT ask questions. Just produce the structured output.
+
+type: object
+properties:
+  result:
+    type: object
+    description: Structured analysis result
+    properties:
+      summary:
+        type: string
+        description: Brief summary of the input
+      sentiment:
+        type: string
+        enum: [positive, negative, neutral]
+        description: Overall sentiment
+      keywords:
+        type: array
+        items:
+          type: string
+        description: Key words extracted from input
+    required: [summary, sentiment, keywords]
+    additionalProperties: false
+
+required:
+  - result
+
+json_schema_extra:
+  kind: agent
+  name: test_structured_output
+  version: "1.0.0"
+  tags: [test, structured-output]
+  structured_output: true
+  tools: []
+  resources: []

rem/services/content/providers.py

@@ -118,15 +118,40 @@ class DocProvider(ContentProvider):
     - Images (.png, .jpg) - OCR text extraction
 
     Handles:
-    - Text extraction with OCR fallback
+    - Text extraction with automatic OCR fallback for scanned documents
     - Table detection and extraction
     - Daemon process workaround for multiprocessing restrictions
+
+    Environment Variables:
+        EXTRACTION_OCR_FALLBACK: Enable OCR fallback (default: true)
+        EXTRACTION_OCR_THRESHOLD: Min chars before triggering OCR fallback (default: 100)
+        EXTRACTION_FORCE_OCR: Always use OCR, skip native extraction (default: false)
+        EXTRACTION_OCR_LANGUAGE: Tesseract language codes (default: eng)
     """
 
     @property
     def name(self) -> str:
         return "doc"
 
+    def _get_env_bool(self, key: str, default: bool) -> bool:
+        """Get boolean from environment variable."""
+        import os
+        val = os.environ.get(key, "").lower()
+        if val in ("true", "1", "yes"):
+            return True
+        elif val in ("false", "0", "no"):
+            return False
+        return default
+
+    def _get_env_int(self, key: str, default: int) -> int:
+        """Get integer from environment variable."""
+        import os
+        val = os.environ.get(key, "")
+        try:
+            return int(val) if val else default
+        except ValueError:
+            return default
+
     def _is_daemon_process(self) -> bool:
         """Check if running in a daemon process."""
         try:
@@ -134,29 +159,34 @@ class DocProvider(ContentProvider):
         except Exception:
             return False
 
-    def _parse_in_subprocess(self, file_path: Path) -> dict:
+    def _parse_in_subprocess(self, file_path: Path, force_ocr: bool = False) -> dict:
         """Run kreuzberg in a separate subprocess to bypass daemon restrictions."""
-        script = """
+        import os
+        ocr_language = os.environ.get("EXTRACTION_OCR_LANGUAGE", "eng")
+
+        script = f"""
 import json
 import sys
 from pathlib import Path
-from kreuzberg import ExtractionConfig, extract_file_sync
+from kreuzberg import ExtractionConfig, OcrConfig, extract_file_sync
 
-# Parse document with kreuzberg 3.x
-config = ExtractionConfig(
-    extract_tables=True,
-    chunk_content=False,
-    extract_keywords=False,
-)
+force_ocr = {force_ocr}
+
+if force_ocr:
+    config = ExtractionConfig(
+        force_ocr=True,
+        ocr=OcrConfig(backend="tesseract", language="{ocr_language}")
+    )
+else:
+    config = ExtractionConfig()
 
 result = extract_file_sync(Path(sys.argv[1]), config=config)
 
-# Serialize result to JSON
-output = {
+output = {{
     'content': result.content,
-    'tables': [t.model_dump() for t in result.tables] if result.tables else [],
-    'metadata': result.metadata
-}
+    'tables': [],
+    'metadata': {{}}
+}}
 print(json.dumps(output))
 """
 
@@ -173,9 +203,41 @@ print(json.dumps(output))
 
         return json.loads(result.stdout)
 
+    def _extract_with_config(self, tmp_path: Path, force_ocr: bool = False) -> tuple[str, dict]:
+        """Extract content with optional OCR config."""
+        import os
+        from kreuzberg import ExtractionConfig, OcrConfig, extract_file_sync
+
+        ocr_language = os.environ.get("EXTRACTION_OCR_LANGUAGE", "eng")
+
+        if force_ocr:
+            config = ExtractionConfig(
+                force_ocr=True,
+                ocr=OcrConfig(backend="tesseract", language=ocr_language)
+            )
+            parser_name = "kreuzberg_ocr"
+        else:
+            config = ExtractionConfig()
+            parser_name = "kreuzberg"
+
+        result = extract_file_sync(tmp_path, config=config)
+        text = result.content
+
+        extraction_metadata = {
+            "parser": parser_name,
+            "file_extension": tmp_path.suffix,
+        }
+
+        return text, extraction_metadata
+
     def extract(self, content: bytes, metadata: dict[str, Any]) -> dict[str, Any]:
         """
-        Extract document content using Kreuzberg.
+        Extract document content using Kreuzberg with intelligent OCR fallback.
+
+        Process:
+        1. Try native text extraction first (fast, preserves structure)
+        2. If content is minimal (< threshold chars), retry with OCR
+        3. Use OCR result if it's better than native result
 
         Args:
             content: Document file bytes
@@ -184,49 +246,89 @@ print(json.dumps(output))
         Returns:
             dict with text and extraction metadata
         """
+        # Get OCR settings from environment
+        force_ocr = self._get_env_bool("EXTRACTION_FORCE_OCR", False)
+        ocr_fallback = self._get_env_bool("EXTRACTION_OCR_FALLBACK", True)
+        ocr_threshold = self._get_env_int("EXTRACTION_OCR_THRESHOLD", 100)
+
         # Write bytes to temp file for kreuzberg
-        # Detect extension from metadata
         content_type = metadata.get("content_type", "")
         suffix = get_extension(content_type, default=".pdf")
 
         with temp_file_from_bytes(content, suffix=suffix) as tmp_path:
+            ocr_used = False
+            ocr_fallback_triggered = False
+            native_char_count = 0
+
             # Check if running in daemon process
             if self._is_daemon_process():
-                logger.info("Daemon process detected - using subprocess workaround for document parsing")
+                logger.info("Daemon process detected - using subprocess workaround")
                 try:
-                    result_dict = self._parse_in_subprocess(tmp_path)
-                    text = result_dict["content"]
-                    extraction_metadata = {
-                        "table_count": len(result_dict["tables"]),
-                        "parser": "kreuzberg_subprocess",
-                        "file_extension": tmp_path.suffix,
-                    }
+                    if force_ocr:
+                        result_dict = self._parse_in_subprocess(tmp_path, force_ocr=True)
+                        text = result_dict["content"]
+                        ocr_used = True
+                        extraction_metadata = {
+                            "parser": "kreuzberg_subprocess_ocr",
+                            "file_extension": tmp_path.suffix,
+                        }
+                    else:
+                        # Try native first
+                        result_dict = self._parse_in_subprocess(tmp_path, force_ocr=False)
+                        text = result_dict["content"]
+                        native_char_count = len(text)
+
+                        # OCR fallback if content is minimal
+                        if ocr_fallback and len(text.strip()) < ocr_threshold:
+                            logger.warning(f"Content below threshold ({len(text.strip())} < {ocr_threshold}) - trying OCR fallback")
+                            try:
+                                ocr_result = self._parse_in_subprocess(tmp_path, force_ocr=True)
+                                ocr_text = ocr_result["content"]
+                                if len(ocr_text.strip()) > len(text.strip()):
+                                    logger.info(f"OCR fallback improved result: {len(ocr_text)} chars (was {native_char_count})")
+                                    text = ocr_text
+                                    ocr_used = True
+                                    ocr_fallback_triggered = True
+                            except Exception as e:
+                                logger.warning(f"OCR fallback failed in subprocess: {e}")
+
+                        extraction_metadata = {
+                            "parser": "kreuzberg_subprocess" if not ocr_used else "kreuzberg_subprocess_ocr_fallback",
+                            "file_extension": tmp_path.suffix,
+                        }
                 except Exception as e:
-                    logger.error(f"Subprocess parsing failed: {e}. Falling back to text-only.")
-                    # Fallback to simple text extraction (kreuzberg 3.x API)
-                    from kreuzberg import ExtractionConfig, extract_file_sync
-                    config = ExtractionConfig(extract_tables=False)
-                    result = extract_file_sync(tmp_path, config=config)
-                    text = result.content
-                    extraction_metadata = {
-                        "parser": "kreuzberg_fallback",
-                        "file_extension": tmp_path.suffix,
-                    }
+                    logger.error(f"Subprocess parsing failed: {e}. Falling back to direct call.")
+                    text, extraction_metadata = self._extract_with_config(tmp_path, force_ocr=force_ocr)
+                    ocr_used = force_ocr
             else:
-                # Normal execution (not in daemon) - kreuzberg 4.x with native ONNX/Rust
-                from kreuzberg import ExtractionConfig, extract_file_sync
-                config = ExtractionConfig(
-                    enable_quality_processing=True,  # Enables table extraction with native ONNX
-                    chunk_content=False,  # We handle chunking ourselves
-                    extract_tables=False,  # Disable table extraction to avoid PyTorch dependency
-                )
-                result = extract_file_sync(tmp_path, config=config)
-                text = result.content
-                extraction_metadata = {
-                    "table_count": len(result.tables) if result.tables else 0,
-                    "parser": "kreuzberg",
-                    "file_extension": tmp_path.suffix,
-                }
+                # Normal execution (not in daemon)
+                if force_ocr:
+                    text, extraction_metadata = self._extract_with_config(tmp_path, force_ocr=True)
+                    ocr_used = True
+                else:
+                    # Try native first
+                    text, extraction_metadata = self._extract_with_config(tmp_path, force_ocr=False)
+                    native_char_count = len(text)
+
+                    # OCR fallback if content is minimal
+                    if ocr_fallback and len(text.strip()) < ocr_threshold:
+                        logger.warning(f"Content below threshold ({len(text.strip())} < {ocr_threshold}) - trying OCR fallback")
+                        try:
+                            ocr_text, _ = self._extract_with_config(tmp_path, force_ocr=True)
+                            if len(ocr_text.strip()) > len(text.strip()):
+                                logger.info(f"OCR fallback improved result: {len(ocr_text)} chars (was {native_char_count})")
+                                text = ocr_text
+                                ocr_used = True
+                                ocr_fallback_triggered = True
+                                extraction_metadata["parser"] = "kreuzberg_ocr_fallback"
+                        except Exception as e:
+                            logger.warning(f"OCR fallback failed: {e}")
+
+            # Add OCR metadata
+            extraction_metadata["ocr_used"] = ocr_used
+            extraction_metadata["ocr_fallback_triggered"] = ocr_fallback_triggered
+            extraction_metadata["native_char_count"] = native_char_count
+            extraction_metadata["final_char_count"] = len(text)
 
             return {
                 "text": text,

rem/services/content/service.py

@@ -274,7 +274,7 @@ class ContentService:
     async def ingest_file(
         self,
         file_uri: str,
-        user_id: str,
+        user_id: str | None = None,
         category: str | None = None,
         tags: list[str] | None = None,
         is_local_server: bool = False,
@@ -283,6 +283,10 @@ class ContentService:
         """
         Complete file ingestion pipeline: read → store → parse → chunk → embed.
 
+        **IMPORTANT: Data is PUBLIC by default (user_id=None).**
+        This is correct for shared knowledge bases (ontologies, procedures, reference data).
+        Private user-scoped data is rarely needed - only set user_id for truly personal content.
+
         **CENTRALIZED INGESTION**: This is the single entry point for all file ingestion
         in REM. It handles:
 
@@ -319,7 +323,9 @@ class ContentService:
 
         Args:
             file_uri: Source file location (local path, s3://, or https://)
-            user_id: User identifier for data isolation and ownership
+            user_id: User identifier for PRIVATE data only. Default None = PUBLIC/shared.
+                Leave as None for shared knowledge bases, ontologies, reference data.
+                Only set for truly private user-specific content.
             category: Optional category tag (document, code, audio, etc.)
             tags: Optional list of tags
             is_local_server: True if running as local/stdio MCP server
@@ -347,12 +353,19 @@ class ContentService:
 
         Example:
             >>> service = ContentService()
+            >>> # PUBLIC data (default) - visible to all users
             >>> result = await service.ingest_file(
-            ...     file_uri="s3://bucket/contract.pdf",
-            ...     user_id="user-123",
-            ...     category="legal"
+            ...     file_uri="s3://bucket/procedure.pdf",
+            ...     category="medical"
             ... )
             >>> print(f"Created {result['resources_created']} searchable chunks")
+            >>>
+            >>> # PRIVATE data (rare) - only for user-specific content
+            >>> result = await service.ingest_file(
+            ...     file_uri="s3://bucket/personal-notes.pdf",
+            ...     user_id="user-123",  # Only this user can access
+            ...     category="personal"
+            ... )
         """
         from pathlib import Path
         from uuid import uuid4

rem/services/embeddings/worker.py

@@ -23,6 +23,8 @@ Future:
 import asyncio
 import os
 from typing import Any, Optional
+import hashlib
+import uuid
 from uuid import uuid4
 
 import httpx
@@ -108,6 +110,7 @@ class EmbeddingWorker:
         self.task_queue: asyncio.Queue = asyncio.Queue()
         self.workers: list[asyncio.Task] = []
         self.running = False
+        self._in_flight_count = 0  # Track tasks being processed (not just in queue)
 
         # Store API key for direct HTTP requests
         from ...settings import settings
@@ -143,17 +146,18 @@ class EmbeddingWorker:
             return
 
         queue_size = self.task_queue.qsize()
-        logger.debug(f"Stopping EmbeddingWorker (processing {queue_size} queued tasks first)")
+        in_flight = self._in_flight_count
+        logger.debug(f"Stopping EmbeddingWorker (queue={queue_size}, in_flight={in_flight})")
 
-        # Wait for queue to drain (with timeout)
+        # Wait for both queue to drain AND in-flight tasks to complete
         max_wait = 30  # 30 seconds max
         waited = 0.0
-        while not self.task_queue.empty() and waited < max_wait:
+        while (not self.task_queue.empty() or self._in_flight_count > 0) and waited < max_wait:
             await asyncio.sleep(0.5)
             waited += 0.5
 
-        if not self.task_queue.empty():
-            remaining = self.task_queue.qsize()
+        if not self.task_queue.empty() or self._in_flight_count > 0:
+            remaining = self.task_queue.qsize() + self._in_flight_count
             logger.warning(
                 f"EmbeddingWorker timeout: {remaining} tasks remaining after {max_wait}s"
             )
@@ -205,12 +209,18 @@ class EmbeddingWorker:
                 if not batch:
                     continue
 
-                logger.debug(f"Worker {worker_id} processing batch of {len(batch)} tasks")
+                # Track in-flight tasks
+                self._in_flight_count += len(batch)
 
-                # Generate embeddings for batch
-                await self._process_batch(batch)
+                logger.debug(f"Worker {worker_id} processing batch of {len(batch)} tasks")
 
-                logger.debug(f"Worker {worker_id} completed batch")
+                try:
+                    # Generate embeddings for batch
+                    await self._process_batch(batch)
+                    logger.debug(f"Worker {worker_id} completed batch")
+                finally:
+                    # Always decrement in-flight count, even on error
+                    self._in_flight_count -= len(batch)
 
             except asyncio.CancelledError:
                 logger.debug(f"Worker {worker_id} cancelled")
@@ -373,7 +383,11 @@ class EmbeddingWorker:
         for task, embedding in zip(tasks, embeddings):
             table_name = f"embeddings_{task.table_name}"
 
-            # Build upsert SQL
+            # Generate deterministic ID from key fields (entity_id, field_name, provider)
+            key_string = f"{task.entity_id}:{task.field_name}:{task.provider}"
+            embedding_id = str(uuid.UUID(hashlib.md5(key_string.encode()).hexdigest()))
+
+            # Build upsert SQL - conflict on deterministic ID
             sql = f"""
                 INSERT INTO {table_name} (
                     id,
@@ -386,7 +400,7 @@ class EmbeddingWorker:
                     updated_at
                 )
                 VALUES ($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
-                ON CONFLICT (entity_id, field_name, provider)
+                ON CONFLICT (id)
                 DO UPDATE SET
                     model = EXCLUDED.model,
                     embedding = EXCLUDED.embedding,
@@ -400,7 +414,7 @@ class EmbeddingWorker:
                 await self.postgres_service.execute(
                     sql,
                     (
-                        str(uuid4()),
+                        embedding_id,
                         task.entity_id,
                         task.field_name,
                         task.provider,

rem/services/postgres/__init__.py

@@ -3,22 +3,47 @@ PostgreSQL service for CloudNativePG database operations.
 """
 
 from .diff_service import DiffService, SchemaDiff
+from .programmable_diff_service import (
+    DiffResult,
+    ObjectDiff,
+    ObjectType,
+    ProgrammableDiffService,
+)
 from .repository import Repository
 from .service import PostgresService
 
 
+_postgres_instance: PostgresService | None = None
+
+
 def get_postgres_service() -> PostgresService | None:
     """
-    Get PostgresService instance.
+    Get PostgresService singleton instance.
 
     Returns None if Postgres is disabled.
+    Uses singleton pattern to prevent connection pool exhaustion.
     """
+    global _postgres_instance
+
     from ...settings import settings
 
     if not settings.postgres.enabled:
         return None
 
-    return PostgresService()
+    if _postgres_instance is None:
+        _postgres_instance = PostgresService()
+
+    return _postgres_instance
 
 
-__all__ = ["PostgresService", "get_postgres_service", "Repository", "DiffService", "SchemaDiff"]
+__all__ = [
+    "DiffResult",
+    "DiffService",
+    "ObjectDiff",
+    "ObjectType",
+    "PostgresService",
+    "ProgrammableDiffService",
+    "Repository",
+    "SchemaDiff",
+    "get_postgres_service",
+]

rem/services/postgres/diff_service.py

@@ -5,12 +5,17 @@ Uses Alembic autogenerate to detect differences between:
 - Target schema (derived from Pydantic models)
 - Current database schema
 
+Also compares programmable objects (functions, triggers, views) which
+Alembic does not track.
+
 This enables:
 1. Local development: See what would change before applying migrations
 2. CI validation: Detect drift between code and database (--check mode)
 3. Migration generation: Create incremental migration files
 """
 
+import asyncio
+import re
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Optional
@@ -51,11 +56,14 @@ class SchemaDiff:
     sql: str = ""
     upgrade_ops: Optional[ops.UpgradeOps] = None
     filtered_count: int = 0  # Number of operations filtered out by strategy
+    # Programmable objects (functions, triggers, views)
+    programmable_summary: list[str] = field(default_factory=list)
+    programmable_sql: str = ""
 
     @property
     def change_count(self) -> int:
         """Total number of detected changes."""
-        return len(self.summary)
+        return len(self.summary) + len(self.programmable_summary)
 
 
 class DiffService:
@@ -127,10 +135,13 @@ class DiffService:
             # These are now generated in pydantic_to_sqlalchemy
         return True
 
-    def compute_diff(self) -> SchemaDiff:
+    def compute_diff(self, include_programmable: bool = True) -> SchemaDiff:
         """
         Compare Pydantic models against database and return differences.
 
+        Args:
+            include_programmable: If True, also diff functions/triggers/views
+
         Returns:
             SchemaDiff with detected changes
         """
@@ -167,21 +178,62 @@ class DiffService:
                 for op in filtered_ops:
                     summary.extend(self._describe_operation(op))
 
-        has_changes = len(summary) > 0
-
         # Generate SQL if there are changes
         sql = ""
-        if has_changes and upgrade_ops:
+        if summary and upgrade_ops:
             sql = self._render_sql(upgrade_ops, engine)
 
+        # Programmable objects diff (functions, triggers, views)
+        programmable_summary = []
+        programmable_sql = ""
+        if include_programmable:
+            prog_summary, prog_sql = self._compute_programmable_diff()
+            programmable_summary = prog_summary
+            programmable_sql = prog_sql
+
+        has_changes = len(summary) > 0 or len(programmable_summary) > 0
+
         return SchemaDiff(
             has_changes=has_changes,
             summary=summary,
             sql=sql,
             upgrade_ops=upgrade_ops,
             filtered_count=filtered_count,
+            programmable_summary=programmable_summary,
+            programmable_sql=programmable_sql,
         )
 
+    def _compute_programmable_diff(self) -> tuple[list[str], str]:
+        """
+        Compute diff for programmable objects (functions, triggers, views).
+
+        Returns:
+            Tuple of (summary_lines, sync_sql)
+        """
+        from .programmable_diff_service import ProgrammableDiffService
+
+        service = ProgrammableDiffService()
+
+        # Run async diff in sync context
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        result = loop.run_until_complete(service.compute_diff())
+
+        summary = []
+        for diff in result.diffs:
+            if diff.status == "missing":
+                summary.append(f"+ {diff.object_type.value.upper()} {diff.name} (missing)")
+            elif diff.status == "different":
+                summary.append(f"~ {diff.object_type.value.upper()} {diff.name} (different)")
+            elif diff.status == "extra":
+                summary.append(f"- {diff.object_type.value.upper()} {diff.name} (extra in db)")
+
+        return summary, result.sync_sql
+
     def _filter_operations(self, operations: list) -> tuple[list, int]:
         """
         Filter operations based on migration strategy.