remdb 0.3.0__py3-none-any.whl → 0.3.114__py3-none-any.whl

rem/services/embeddings/api.py

@@ -5,13 +5,20 @@ Provides synchronous and async wrappers for embedding generation using
 raw HTTP requests (no OpenAI SDK dependency).
 """
 
-import os
 from typing import Optional, cast
 
 import httpx
 import requests
 from loguru import logger
 
+from rem.utils.constants import DEFAULT_EMBEDDING_DIMS, HTTP_TIMEOUT_DEFAULT
+
+
+def _get_openai_api_key() -> Optional[str]:
+    """Get OpenAI API key from settings."""
+    from rem.settings import settings
+    return settings.llm.openai_api_key
+
 
 def generate_embedding(
     text: str,
@@ -26,19 +33,19 @@ def generate_embedding(
         text: Text to embed
         model: Model name (default: text-embedding-3-small)
         provider: Provider name (default: openai)
-        api_key: API key (defaults to OPENAI_API_KEY env var)
+        api_key: API key (defaults to settings.llm.openai_api_key)
 
     Returns:
         Embedding vector (1536 dimensions for text-embedding-3-small)
     """
     if provider == "openai":
-        api_key = api_key or os.getenv("OPENAI_API_KEY")
+        api_key = api_key or _get_openai_api_key()
         if not api_key:
             logger.warning("No OpenAI API key - returning zero vector")
-            return [0.0] * 1536
+            return [0.0] * DEFAULT_EMBEDDING_DIMS
 
         try:
-            logger.info(f"Generating OpenAI embedding for text using {model}")
+            logger.debug(f"Generating OpenAI embedding for text using {model}")
 
             response = requests.post(
                 "https://api.openai.com/v1/embeddings",
@@ -47,22 +54,22 @@ def generate_embedding(
                     "Content-Type": "application/json",
                 },
                 json={"input": [text], "model": model},
-                timeout=30,
+                timeout=HTTP_TIMEOUT_DEFAULT,
             )
             response.raise_for_status()
 
             data = response.json()
             embedding = data["data"][0]["embedding"]
-            logger.info(f"Successfully generated embedding (dimension: {len(embedding)})")
+            logger.debug(f"Successfully generated embedding (dimension: {len(embedding)})")
             return cast(list[float], embedding)
 
         except Exception as e:
             logger.error(f"Failed to generate embedding from OpenAI: {e}", exc_info=True)
-            return [0.0] * 1536
+            return [0.0] * DEFAULT_EMBEDDING_DIMS
 
     else:
         logger.warning(f"Unsupported provider '{provider}' - returning zero vector")
-        return [0.0] * 1536
+        return [0.0] * DEFAULT_EMBEDDING_DIMS
 
 
 async def generate_embedding_async(
@@ -78,19 +85,19 @@ async def generate_embedding_async(
         text: Text to embed
         model: Model name (default: text-embedding-3-small)
         provider: Provider name (default: openai)
-        api_key: API key (defaults to OPENAI_API_KEY env var)
+        api_key: API key (defaults to settings.llm.openai_api_key)
 
     Returns:
         Embedding vector (1536 dimensions for text-embedding-3-small)
     """
     if provider == "openai":
-        api_key = api_key or os.getenv("OPENAI_API_KEY")
+        api_key = api_key or _get_openai_api_key()
         if not api_key:
             logger.warning("No OpenAI API key - returning zero vector")
-            return [0.0] * 1536
+            return [0.0] * DEFAULT_EMBEDDING_DIMS
 
         try:
-            logger.info(f"Generating OpenAI embedding for text using {model}")
+            logger.debug(f"Generating OpenAI embedding for text using {model}")
 
             async with httpx.AsyncClient() as client:
                 response = await client.post(
@@ -100,21 +107,21 @@ async def generate_embedding_async(
                         "Content-Type": "application/json",
                     },
                     json={"input": [text], "model": model},
-                    timeout=30.0,
+                    timeout=HTTP_TIMEOUT_DEFAULT,
                 )
                 response.raise_for_status()
 
                 data = response.json()
                 embedding = data["data"][0]["embedding"]
-                logger.info(
+                logger.debug(
                     f"Successfully generated embedding (dimension: {len(embedding)})"
                 )
                 return cast(list[float], embedding)
 
         except Exception as e:
             logger.error(f"Failed to generate embedding from OpenAI: {e}", exc_info=True)
-            return [0.0] * 1536
+            return [0.0] * DEFAULT_EMBEDDING_DIMS
 
     else:
         logger.warning(f"Unsupported provider '{provider}' - returning zero vector")
-        return [0.0] * 1536
+        return [0.0] * DEFAULT_EMBEDDING_DIMS

rem/services/embeddings/worker.py

@@ -69,7 +69,7 @@ def get_global_embedding_worker(postgres_service: Any = None) -> "EmbeddingWorke
         if postgres_service is None:
             raise RuntimeError("Must provide postgres_service on first call to get_global_embedding_worker")
         _global_worker = EmbeddingWorker(postgres_service=postgres_service)
-        logger.info("Created global EmbeddingWorker singleton")
+        logger.debug("Created global EmbeddingWorker singleton")
 
     return _global_worker
 
@@ -117,7 +117,7 @@ class EmbeddingWorker:
                 "No OpenAI API key provided - embeddings will use zero vectors"
             )
 
-        logger.info(
+        logger.debug(
             f"Initialized EmbeddingWorker: {num_workers} workers, "
             f"batch_size={batch_size}, timeout={batch_timeout}s"
         )
@@ -125,17 +125,17 @@ class EmbeddingWorker:
     async def start(self) -> None:
         """Start worker pool."""
         if self.running:
-            logger.warning("EmbeddingWorker already running")
+            logger.debug("EmbeddingWorker already running")
             return
 
         self.running = True
-        logger.info(f"Starting {self.num_workers} embedding workers")
+        logger.debug(f"Starting {self.num_workers} embedding workers")
 
         for i in range(self.num_workers):
             worker = asyncio.create_task(self._worker_loop(i))
             self.workers.append(worker)
 
-        logger.info("EmbeddingWorker started")
+        logger.debug("EmbeddingWorker started")
 
     async def stop(self) -> None:
         """Stop worker pool gracefully - processes remaining queue before stopping."""
@@ -143,7 +143,7 @@ class EmbeddingWorker:
             return
 
         queue_size = self.task_queue.qsize()
-        logger.info(f"Stopping EmbeddingWorker (processing {queue_size} queued tasks first)")
+        logger.debug(f"Stopping EmbeddingWorker (processing {queue_size} queued tasks first)")
 
         # Wait for queue to drain (with timeout)
         max_wait = 30  # 30 seconds max
@@ -171,7 +171,7 @@ class EmbeddingWorker:
         await asyncio.gather(*self.workers, return_exceptions=True)
 
         self.workers.clear()
-        logger.info("EmbeddingWorker stopped")
+        logger.debug("EmbeddingWorker stopped")
 
     async def queue_task(self, task: EmbeddingTask) -> None:
         """
@@ -195,7 +195,7 @@ class EmbeddingWorker:
         Args:
             worker_id: Unique worker identifier
         """
-        logger.info(f"Worker {worker_id} started")
+        logger.debug(f"Worker {worker_id} started")
 
         while self.running:
             try:
@@ -205,7 +205,7 @@ class EmbeddingWorker:
                 if not batch:
                     continue
 
-                logger.info(f"Worker {worker_id} processing batch of {len(batch)} tasks")
+                logger.debug(f"Worker {worker_id} processing batch of {len(batch)} tasks")
 
                 # Generate embeddings for batch
                 await self._process_batch(batch)
@@ -213,14 +213,14 @@ class EmbeddingWorker:
                 logger.debug(f"Worker {worker_id} completed batch")
 
             except asyncio.CancelledError:
-                logger.info(f"Worker {worker_id} cancelled")
+                logger.debug(f"Worker {worker_id} cancelled")
                 break
             except Exception as e:
                 logger.error(f"Worker {worker_id} error: {e}", exc_info=True)
                 # Continue processing (don't crash worker on error)
                 await asyncio.sleep(1)
 
-        logger.info(f"Worker {worker_id} stopped")
+        logger.debug(f"Worker {worker_id} stopped")
 
     async def _collect_batch(self) -> list[EmbeddingTask]:
         """
@@ -284,10 +284,10 @@ class EmbeddingWorker:
             )
 
             # Upsert to database
-            logger.info(f"Upserting {len(embeddings)} embeddings to database...")
+            logger.debug(f"Upserting {len(embeddings)} embeddings to database...")
             await self._upsert_embeddings(batch, embeddings)
 
-            logger.info(
+            logger.debug(
                 f"Successfully generated and stored {len(embeddings)} embeddings "
                 f"(provider={provider}, model={model})"
             )
@@ -315,7 +315,7 @@ class EmbeddingWorker:
         """
         if provider == "openai" and self.openai_api_key:
             try:
-                logger.info(
+                logger.debug(
                     f"Generating OpenAI embeddings for {len(texts)} texts using {model}"
                 )
 
@@ -336,7 +336,7 @@ class EmbeddingWorker:
                     data = response.json()
                     embeddings = [item["embedding"] for item in data["data"]]
 
-                    logger.info(
+                    logger.debug(
                         f"Successfully generated {len(embeddings)} embeddings from OpenAI"
                     )
                     return embeddings
@@ -409,7 +409,7 @@ class EmbeddingWorker:
                     ),
                 )
 
-                logger.info(
+                logger.debug(
                     f"Upserted embedding: {task.table_name}.{task.entity_id}.{task.field_name}"
                 )
 

rem/services/phoenix/EXPERIMENT_DESIGN.md

@@ -164,7 +164,7 @@ cp curated-queries.csv experiments/rem-001/validation/production/
 **Option C: Curated Engrams**
 ```bash
 # Generate engrams from REM data
-rem dreaming full --user-id test-user --tenant-id acme --generate-test-cases
+rem dreaming full --user-id test-user  --generate-test-cases
 
 # Review and select high-quality engrams
 rem engram list --quality high --limit 100 --output engrams.csv
@@ -357,7 +357,7 @@ Level 4 (Mature):   Multiple cycles, full query capabilities
 # Generate engrams from REM data
 rem dreaming full \
   --user-id test-user \
-  --tenant-id acme \
+   \
   --generate-test-cases \
   --quality-level 3
 
@@ -1027,7 +1027,7 @@ rem experiments experiment run rem-lookup-ask_rem-golden \
 
 ```bash
 # 1. Generate high-quality engrams
-rem dreaming full --tenant-id acme --generate-test-cases --quality-level 4
+rem dreaming full  --generate-test-cases --quality-level 4
 
 # 2. Export engrams
 rem engram export rem-engrams-mature-mixed --output engrams.csv --format phoenix

rem/services/phoenix/client.py

@@ -53,7 +53,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any, Callable, TYPE_CHECKING, cast
 
-import pandas as pd
+import polars as pl
 from loguru import logger
 
 from .config import PhoenixConfig
@@ -64,6 +64,95 @@ if TYPE_CHECKING:
     from phoenix.client.resources.experiments.types import RanExperiment
 
 
+def dataframe_to_phoenix_dataset(
+    client: "PhoenixClient",
+    df: pl.DataFrame,
+    dataset_name: str,
+    input_keys: list[str] | None = None,
+    output_keys: list[str] | None = None,
+    metadata_keys: list[str] | None = None,
+    description: str | None = None,
+) -> "Dataset":
+    """Convert a Polars DataFrame to a Phoenix Dataset.
+
+    This function transforms a Polars DataFrame into a Phoenix Dataset by:
+    1. Extracting input columns (what agents receive)
+    2. Extracting output columns (ground truth/expected output)
+    3. Extracting metadata columns (optional labels, difficulty, etc.)
+
+    If column keys are not specified, uses smart defaults:
+    - input_keys: columns containing 'input', 'query', 'question', or 'prompt'
+    - output_keys: columns containing 'output', 'expected', 'answer', or 'response'
+    - metadata_keys: remaining columns
+
+    Args:
+        client: PhoenixClient instance
+        df: Polars DataFrame with experiment data
+        dataset_name: Name for the created Phoenix dataset
+        input_keys: Optional list of column names for inputs
+        output_keys: Optional list of column names for outputs (ground truth)
+        metadata_keys: Optional list of column names for metadata
+        description: Optional dataset description
+
+    Returns:
+        Phoenix Dataset instance
+
+    Example:
+        >>> df = pl.read_csv("golden_set.csv")
+        >>> dataset = dataframe_to_phoenix_dataset(
+        ...     client=phoenix_client,
+        ...     df=df,
+        ...     dataset_name="my-golden-set",
+        ...     input_keys=["query"],
+        ...     output_keys=["expected_output"],
+        ...     metadata_keys=["difficulty"]
+        ... )
+    """
+    columns = df.columns
+
+    # Smart defaults for column detection
+    if input_keys is None:
+        input_keys = [c for c in columns if any(
+            k in c.lower() for k in ["input", "query", "question", "prompt"]
+        )]
+        if not input_keys:
+            # Fallback: first column
+            input_keys = [columns[0]] if columns else []
+
+    if output_keys is None:
+        output_keys = [c for c in columns if any(
+            k in c.lower() for k in ["output", "expected", "answer", "response", "reference"]
+        )]
+        if not output_keys:
+            # Fallback: second column
+            output_keys = [columns[1]] if len(columns) > 1 else []
+
+    if metadata_keys is None:
+        used_keys = set(input_keys) | set(output_keys)
+        metadata_keys = [c for c in columns if c not in used_keys]
+
+    logger.debug(
+        f"DataFrame to Phoenix Dataset: inputs={input_keys}, "
+        f"outputs={output_keys}, metadata={metadata_keys}"
+    )
+
+    # Convert to list of dicts
+    records = df.to_dicts()
+
+    inputs = [{k: row.get(k) for k in input_keys} for row in records]
+    outputs = [{k: row.get(k) for k in output_keys} for row in records]
+    metadata = [{k: row.get(k) for k in metadata_keys} for row in records] if metadata_keys else None
+
+    # Create Phoenix dataset
+    return client.create_dataset_from_data(
+        name=dataset_name,
+        inputs=inputs,
+        outputs=outputs,
+        metadata=metadata,
+        description=description,
+    )
+
+
 class PhoenixClient:
     """High-level Phoenix client for REM evaluation workflows.
 
@@ -260,19 +349,22 @@ class PhoenixClient:
             "SEARCH semantic AI engineer",sarah-chen,person,medium,SEARCH
         """
         try:
-            # Load CSV
-            df = pd.read_csv(csv_file_path)
+            # Load CSV with Polars
+            df = pl.read_csv(csv_file_path)
+
+            # Convert to list of dicts
+            records = df.to_dicts()
 
             # Extract inputs
-            inputs = cast(list[dict[str, Any]], df[input_keys].to_dict("records"))
+            inputs = [{k: row.get(k) for k in input_keys} for row in records]
 
             # Extract outputs
-            outputs = cast(list[dict[str, Any]], df[output_keys].to_dict("records"))
+            outputs = [{k: row.get(k) for k in output_keys} for row in records]
 
             # Extract metadata if specified
             metadata = None
             if metadata_keys:
-                metadata = cast(list[dict[str, Any]], df[metadata_keys].to_dict("records"))
+                metadata = [{k: row.get(k) for k in metadata_keys} for row in records]
 
             return self.create_dataset_from_data(
                 name=name,
@@ -331,13 +423,16 @@ class PhoenixClient:
 
     def run_experiment(
         self,
-        dataset: "Dataset" | str,
+        dataset: "Dataset" | str | pl.DataFrame,
         task: Callable[[Any], Any] | None = None,
         evaluators: list[Callable[[Any], Any]] | None = None,
         experiment_name: str | None = None,
         experiment_description: str | None = None,
         experiment_metadata: dict[str, Any] | None = None,
         experiment_config: Any | None = None,
+        input_keys: list[str] | None = None,
+        output_keys: list[str] | None = None,
+        metadata_keys: list[str] | None = None,
     ) -> "RanExperiment":
         """Run an evaluation experiment.
 
@@ -346,14 +441,22 @@ class PhoenixClient:
         2. Agent run: Provide task function to execute agents on dataset
         3. Evaluator run: Provide evaluators to score existing outputs
 
+        Dataset can be:
+        - Phoenix Dataset instance
+        - Dataset name (string) - will be loaded from Phoenix
+        - Polars DataFrame - will be converted to Phoenix Dataset
+
         Args:
-            dataset: Dataset instance or name (required unless experiment_config provided)
+            dataset: Dataset instance, name, or Polars DataFrame
             task: Optional task function to run on each example (agent execution)
             evaluators: Optional list of evaluator functions
             experiment_name: Optional experiment name
             experiment_description: Optional description
             experiment_metadata: Optional metadata dict
             experiment_config: Optional ExperimentConfig instance (overrides other params)
+            input_keys: Column names for inputs (required if dataset is DataFrame)
+            output_keys: Column names for outputs (required if dataset is DataFrame)
+            metadata_keys: Optional column names for metadata
 
         Returns:
             RanExperiment with results
@@ -369,6 +472,16 @@ class PhoenixClient:
             ...     experiment_name="rem-v1-baseline"
             ... )
 
+        Example - With Polars DataFrame:
+            >>> df = pl.read_csv("golden_set.csv")
+            >>> experiment = client.run_experiment(
+            ...     dataset=df,
+            ...     task=run_agent,
+            ...     experiment_name="rem-v1-baseline",
+            ...     input_keys=["query"],
+            ...     output_keys=["expected_output"]
+            ... )
+
         Example - Evaluator Run (Phase 2b):
             >>> experiment = client.run_experiment(
             ...     dataset=agent_results,
@@ -407,6 +520,21 @@ class PhoenixClient:
                     else:
                         dataset = dataset_ref.path
 
+            # Convert Polars DataFrame to Phoenix Dataset
+            if isinstance(dataset, pl.DataFrame):
+                dataset_name_for_phoenix = f"{experiment_name or 'experiment'}-dataset-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+                logger.info(f"Converting Polars DataFrame to Phoenix Dataset: {dataset_name_for_phoenix}")
+                dataset = dataframe_to_phoenix_dataset(
+                    client=self,
+                    df=dataset,
+                    dataset_name=dataset_name_for_phoenix,
+                    input_keys=input_keys,
+                    output_keys=output_keys,
+                    metadata_keys=metadata_keys,
+                    description=f"Auto-created from DataFrame for experiment: {experiment_name}",
+                )
+                logger.info(f"✓ Created Phoenix Dataset: {dataset_name_for_phoenix}")
+
             # Load dataset if name provided
             if isinstance(dataset, str):
                 dataset = self.get_dataset(dataset)
@@ -454,7 +582,7 @@ class PhoenixClient:
         root_spans_only: bool = True,
         trace_id: str | None = None,
         span_id: str | None = None,
-    ) -> pd.DataFrame:
+    ) -> pl.DataFrame:
         """Query traces from Phoenix.
 
         Args:
@@ -467,7 +595,7 @@ class PhoenixClient:
             span_id: Filter by specific span ID
 
         Returns:
-            DataFrame with trace data
+            Polars DataFrame with trace data
 
         Example:
             >>> traces = client.get_traces(
@@ -492,8 +620,11 @@ class PhoenixClient:
             if span_id:
                 query_params["span_id"] = span_id
 
-            # Query traces
-            traces_df = self._client.query_spans(limit=limit, **query_params)  # type: ignore[attr-defined]
+            # Query traces (Phoenix returns pandas DataFrame)
+            pandas_df = self._client.query_spans(limit=limit, **query_params)  # type: ignore[attr-defined]
+
+            # Convert pandas to Polars
+            traces_df = pl.from_pandas(pandas_df)
 
             logger.debug(f"Retrieved {len(traces_df)} traces")
             return traces_df
@@ -535,7 +666,7 @@ class PhoenixClient:
             ... )
         """
         try:
-            # Query traces
+            # Query traces (returns Polars DataFrame)
             traces_df = self.get_traces(
                 project_name=project_name,
                 start_time=start_time,
@@ -547,12 +678,15 @@ class PhoenixClient:
             if len(traces_df) == 0:
                 raise ValueError("No traces found matching criteria")
 
+            # Convert to list of dicts for iteration
+            records = traces_df.to_dicts()
+
             # Extract inputs and outputs from traces
             inputs = []
             outputs = []
             metadata = []
 
-            for _, row in traces_df.iterrows():
+            for row in records:
                 # Extract input
                 span_input = row.get("attributes.input")
                 if span_input:
@@ -658,29 +792,128 @@ class PhoenixClient:
         label: str | None = None,
         score: float | None = None,
         explanation: str | None = None,
-    ) -> None:
+        metadata: dict[str, Any] | None = None,
+    ) -> str | None:
         """Add feedback annotation to a span.
 
         Args:
             span_id: Span ID to annotate
-            annotation_name: Name of the annotation (e.g., "correctness")
+            annotation_name: Name of the annotation (e.g., "correctness", "user_feedback")
             annotator_kind: Type of annotator ("HUMAN", "LLM", "CODE")
-            label: Optional label (e.g., "correct", "incorrect")
+            label: Optional label (e.g., "correct", "incorrect", "helpful")
             score: Optional numeric score (0.0-1.0)
             explanation: Optional explanation text
+            metadata: Optional additional metadata dict
+
+        Returns:
+            Annotation ID if successful, None otherwise
         """
         try:
-            self._client.add_span_annotation(  # type: ignore[attr-defined]
+            result = self._client.add_span_annotation(  # type: ignore[attr-defined]
                 span_id=span_id,
                 name=annotation_name,
                 annotator_kind=annotator_kind,
                 label=label,
                 score=score,
                 explanation=explanation,
+                metadata=metadata,
             )
 
-            logger.info(f"Added {annotator_kind} feedback to span {span_id}")
+            annotation_id = getattr(result, "id", None) if result else None
+            logger.info(f"Added {annotator_kind} feedback to span {span_id} -> {annotation_id}")
+
+            return annotation_id
 
         except Exception as e:
             logger.error(f"Failed to add span feedback: {e}")
             raise
+
+    def sync_user_feedback(
+        self,
+        span_id: str,
+        rating: int | None = None,
+        categories: list[str] | None = None,
+        comment: str | None = None,
+        feedback_id: str | None = None,
+    ) -> str | None:
+        """Sync user feedback to Phoenix as a span annotation.
+
+        Convenience method for syncing Feedback entities to Phoenix.
+        Converts REM feedback format to Phoenix annotation format.
+
+        Args:
+            span_id: OTEL span ID to annotate
+            rating: User rating (-1, 1-5 scale)
+            categories: List of feedback categories
+            comment: Free-text comment
+            feedback_id: Optional REM feedback ID for reference
+
+        Returns:
+            Phoenix annotation ID if successful
+
+        Example:
+            >>> client.sync_user_feedback(
+            ...     span_id="abc123",
+            ...     rating=4,
+            ...     categories=["helpful", "accurate"],
+            ...     comment="Great response!"
+            ... )
+        """
+        # Convert rating to 0-1 score
+        score = None
+        if rating is not None:
+            if rating == -1:
+                score = 0.0
+            elif 1 <= rating <= 5:
+                score = rating / 5.0
+
+        # Use primary category as label
+        label = categories[0] if categories else None
+
+        # Build explanation from comment and additional categories
+        explanation = comment
+        if categories and len(categories) > 1:
+            cats_str = ", ".join(categories[1:])
+            if explanation:
+                explanation = f"{explanation} [Categories: {cats_str}]"
+            else:
+                explanation = f"Categories: {cats_str}"
+
+        # Build metadata
+        metadata = {
+            "rating": rating,
+            "categories": categories or [],
+        }
+        if feedback_id:
+            metadata["rem_feedback_id"] = feedback_id
+
+        return self.add_span_feedback(
+            span_id=span_id,
+            annotation_name="user_feedback",
+            annotator_kind="HUMAN",
+            label=label,
+            score=score,
+            explanation=explanation,
+            metadata=metadata,
+        )
+
+    def get_span_annotations(
+        self,
+        span_id: str,
+        annotation_name: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Get annotations for a span.
+
+        Args:
+            span_id: Span ID to query
+            annotation_name: Optional filter by annotation name
+
+        Returns:
+            List of annotation dicts
+
+        TODO: Implement once Phoenix client exposes this method
+        """
+        # TODO: Phoenix client doesn't expose annotation query yet
+        # This is a stub for future implementation
+        logger.warning("get_span_annotations not yet implemented in Phoenix client")
+        return []