uncertainty-engine-types 0.0.2__tar.gz

uncertainty_engine_types-0.0.2/PKG-INFO

@@ -0,0 +1,126 @@
+Metadata-Version: 2.3
+Name: uncertainty-engine-types
+Version: 0.0.2
+Summary: Common type definitions for the Uncertainty Engine
+Author: Freddy Wordingham
+Author-email: freddy@digilab.ai
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: pandas (>=2.2.3,<3.0.0)
+Requires-Dist: pydantic (>=2.10.5,<3.0.0)
+Requires-Dist: typeguard (>=2.13.3,<2.14.0)
+Description-Content-Type: text/markdown
+
+<div style="color: black; background-color: #edf497;">
+    <img src="./assets/images/uncertainty-engine-logo.png">
+</div>
+
+# Types
+
+Common types definitions for the Uncertainty Engine.
+This library should be used by other packages to ensure consistency in the types used across the Uncertainty Engine.
+
+## Overview
+
+### Execution & Error Handling
+
+- **ExecutionError**  
+  Exception raised to indicate execution errors.
+
+### Graph & Node Types
+
+- **Graph**  
+  Represents a collection of nodes and their connections.
+- **NodeElement**  
+  Defines a node with a type and associated inputs.
+- **NodeId**  
+  A unique identifier for nodes.
+- **SourceHandle** & **TargetHandle**  
+  Strings used to reference node connections.
+
+### Node Handles
+
+- **Handle**  
+  Represents a node handle in the format `node.handle` and validates this structure.
+
+### Language Learning Models (LLMs)
+
+- **LLM**  
+  Abstract base class for language learning models.
+- **OpenAILLM**  
+  LLM implementation using OpenAI.
+- **OllamaLLM**  
+  LLM implementation using Ollama.
+- **LLMProvider**  
+  Enum listing supported LLM providers.
+- **LLMManager**  
+  Manages connections to LLMs based on the chosen provider and configuration.
+
+### Messaging
+
+- **Message**  
+  Represents a message with a role and content, used for interactions with LLMs.
+
+### TwinLab Models
+
+- **TwinLabModel**  
+  Represents a model configuration including metadata.
+- **save_model**  
+  Function to persist a model configuration.
+
+### Node Metadata
+
+- **NodeInputInfo**  
+  Describes the properties of a node's input.
+- **NodeOutputInfo**  
+  Describes the properties of a node's output.
+- **NodeInfo**  
+  Aggregates metadata for a node, including inputs and outputs.
+
+### Sensor Design
+
+- **SensorDesigner**  
+  Defines sensor configuration and provides functionality to load sensor data.
+- **save_sensor_designer**  
+  Function to persist a sensor designer configuration.
+
+### SQL Database Types
+
+- **SQLDatabase**  
+  Abstract base class for executing SQL queries.
+- **PostgreSQL**  
+  Implementation of SQLDatabase for PostgreSQL.
+- **SQLKind**  
+  Enum listing supported SQL database types.
+- **SQLManager**  
+  Manages connections and operations for SQL databases.
+
+### Tabular Data
+
+- **TabularData**  
+  Represents CSV-based data and includes functionality to load it into a pandas DataFrame.
+
+### Token Types
+
+- **Token**  
+  Enum representing token types, such as TRAINING and STANDARD.
+
+### Vector Stores
+
+- **VectorStoreConnection**  
+  Abstract base class for vector store operations.
+- **WeaviateVectorStoreConnection**  
+  Implements a connection to a Weaviate vector store.
+- **VectorStoreProvider**  
+  Enum for supported vector store providers.
+- **VectorStoreManager**  
+  Manages connections to vector stores.
+- **get_persistent_vector_store**  
+  Function to establish a persistent connection to a Weaviate vector store.
+- **get_embedding_function**  
+  Retrieves an embedding function based on configuration, supporting both HuggingFace and OpenAI options.
+

uncertainty_engine_types-0.0.2/README.md

@@ -0,0 +1,108 @@
+<div style="color: black; background-color: #edf497;">
+    <img src="./assets/images/uncertainty-engine-logo.png">
+</div>
+
+# Types
+
+Common types definitions for the Uncertainty Engine.
+This library should be used by other packages to ensure consistency in the types used across the Uncertainty Engine.
+
+## Overview
+
+### Execution & Error Handling
+
+- **ExecutionError**  
+  Exception raised to indicate execution errors.
+
+### Graph & Node Types
+
+- **Graph**  
+  Represents a collection of nodes and their connections.
+- **NodeElement**  
+  Defines a node with a type and associated inputs.
+- **NodeId**  
+  A unique identifier for nodes.
+- **SourceHandle** & **TargetHandle**  
+  Strings used to reference node connections.
+
+### Node Handles
+
+- **Handle**  
+  Represents a node handle in the format `node.handle` and validates this structure.
+
+### Language Learning Models (LLMs)
+
+- **LLM**  
+  Abstract base class for language learning models.
+- **OpenAILLM**  
+  LLM implementation using OpenAI.
+- **OllamaLLM**  
+  LLM implementation using Ollama.
+- **LLMProvider**  
+  Enum listing supported LLM providers.
+- **LLMManager**  
+  Manages connections to LLMs based on the chosen provider and configuration.
+
+### Messaging
+
+- **Message**  
+  Represents a message with a role and content, used for interactions with LLMs.
+
+### TwinLab Models
+
+- **TwinLabModel**  
+  Represents a model configuration including metadata.
+- **save_model**  
+  Function to persist a model configuration.
+
+### Node Metadata
+
+- **NodeInputInfo**  
+  Describes the properties of a node's input.
+- **NodeOutputInfo**  
+  Describes the properties of a node's output.
+- **NodeInfo**  
+  Aggregates metadata for a node, including inputs and outputs.
+
+### Sensor Design
+
+- **SensorDesigner**  
+  Defines sensor configuration and provides functionality to load sensor data.
+- **save_sensor_designer**  
+  Function to persist a sensor designer configuration.
+
+### SQL Database Types
+
+- **SQLDatabase**  
+  Abstract base class for executing SQL queries.
+- **PostgreSQL**  
+  Implementation of SQLDatabase for PostgreSQL.
+- **SQLKind**  
+  Enum listing supported SQL database types.
+- **SQLManager**  
+  Manages connections and operations for SQL databases.
+
+### Tabular Data
+
+- **TabularData**  
+  Represents CSV-based data and includes functionality to load it into a pandas DataFrame.
+
+### Token Types
+
+- **Token**  
+  Enum representing token types, such as TRAINING and STANDARD.
+
+### Vector Stores
+
+- **VectorStoreConnection**  
+  Abstract base class for vector store operations.
+- **WeaviateVectorStoreConnection**  
+  Implements a connection to a Weaviate vector store.
+- **VectorStoreProvider**  
+  Enum for supported vector store providers.
+- **VectorStoreManager**  
+  Manages connections to vector stores.
+- **get_persistent_vector_store**  
+  Function to establish a persistent connection to a Weaviate vector store.
+- **get_embedding_function**  
+  Retrieves an embedding function based on configuration, supporting both HuggingFace and OpenAI options.

uncertainty_engine_types-0.0.2/pyproject.toml

@@ -0,0 +1,30 @@
+[project]
+name = "uncertainty-engine-types"
+version = "0.0.2"
+description = "Common type definitions for the Uncertainty Engine"
+authors = [
+    { name = "Freddy Wordingham", email = "freddy@digilab.ai" },
+    { name = "Jasper Cantwell", email = "jasper@digilab.ai" },
+    { name = "Jamie Donald-McCann", email = "jamie.donald-mccann@digilab.ai" },
+]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+pandas = "^2.2.3"
+pydantic = "^2.10.5"
+python = ">=3.10,<4.0"
+typeguard = ">=2.13.3,<2.14.0"
+
+[tool.poetry.group.dev.dependencies]
+black = "^24.10.0"
+flake8 = "^7.1.1"
+mypy = "^1.13.0"
+pylint = "^3.3.2"
+pytest = "^8.3.4"
+pytest-cov = "^6.0.0"
+pytest-mock = "^3.14.0"
+radon = "^6.0.1"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"

uncertainty_engine_types-0.0.2/uncertainty_engine_types/__init__.py

@@ -0,0 +1,42 @@
+from .execution_error import ExecutionError
+from .graph import Graph, NodeElement, NodeId, SourceHandle, TargetHandle
+from .handle import Handle
+from .llm import LLM, LLMManager, LLMProvider
+from .message import Message
+from .model import TwinLabModel, save_model
+from .node_info import NodeInfo, NodeInputInfo, NodeOutputInfo
+from .sensor_designer import SensorDesigner, save_sensor_designer
+from .sql import SQLDatabase, SQLKind, SQLManager
+from .tabular_data import TabularData
+from .token import Token
+from .vector_store import VectorStoreConnection, VectorStoreManager, VectorStoreProvider
+
+__all__ = [
+    "ExecutionError",
+    "Graph",
+    "Handle",
+    "LLM",
+    "LLMManager",
+    "LLMProvider",
+    "Message",
+    "Message",
+    "NodeElement",
+    "NodeId",
+    "NodeInfo",
+    "NodeInputInfo",
+    "NodeOutputInfo",
+    "save_model",
+    "save_sensor_designer",
+    "SensorDesigner",
+    "SourceHandle",
+    "SQLDatabase",
+    "SQLKind",
+    "SQLManager",
+    "TabularData",
+    "TargetHandle",
+    "Token",
+    "TwinLabModel",
+    "VectorStoreConnection",
+    "VectorStoreManager",
+    "VectorStoreProvider",
+]

uncertainty_engine_types-0.0.2/uncertainty_engine_types/execution_error.py

@@ -0,0 +1,3 @@
+class ExecutionError(Exception):
+    def __init__(self, message: str):
+        super().__init__(message)

uncertainty_engine_types-0.0.2/uncertainty_engine_types/graph.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel
+
+from .handle import Handle
+
+
+NodeId = str
+TargetHandle = str
+SourceHandle = str
+
+
+class NodeElement(BaseModel):
+    type: str
+    inputs: dict[TargetHandle, Handle] = {}
+
+
+class Graph(BaseModel):
+    nodes: dict[NodeId, NodeElement]

uncertainty_engine_types-0.0.2/uncertainty_engine_types/handle.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel
+
+
+class Handle(BaseModel):
+    node_name: str
+    node_handle: str
+
+    def __init__(self, handle_str: str):
+        if handle_str.count(".") != 1:
+            raise ValueError(
+                "Handle string must contain exactly one dot ('.') separating node and handle"
+            )
+
+        node_name, node_handle = handle_str.split(".")
+        super().__init__(node_name=node_name, node_handle=node_handle)

uncertainty_engine_types-0.0.2/uncertainty_engine_types/llm.py

@@ -0,0 +1,111 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel
+from typeguard import typechecked
+
+
+@typechecked
+class LLM(ABC):
+
+    @property
+    def temperature(self) -> float:
+        return self._temperature
+
+    @temperature.setter
+    def temperature(self, value: float) -> None:
+        if not 0.0 <= value <= 1.0:
+            raise ValueError("Temperature must be between 0.0 and 1.0")
+        self._temperature = value
+
+    @abstractmethod
+    def get_langchain_llm(self, query: str) -> "LLM":
+        """
+        Return a BaseLLM object from Langchain according to the provider type.
+        """
+        pass
+
+    def run_query(self, query: str) -> str:
+        """
+        Call the LLM with a query.
+        """
+        llm = self.get_langchain_llm()
+        return llm.invoke(query).content
+
+
+@typechecked
+class OpenAILLM(LLM):
+
+    def __init__(self, api_key: str, model: str, temperature: float = 0.0):
+        self.api_key = api_key
+        self.model = model
+        self.temperature = temperature
+
+    def get_langchain_llm(self):
+        from langchain_openai import ChatOpenAI
+
+        return ChatOpenAI(
+            model=self.model,
+            temperature=self.temperature,
+            api_key=self.api_key,
+        )
+
+
+@typechecked
+class OllamaLLM(LLM):
+    def __init__(self, url: str, model: str, temperature: float = 0.0):
+        self.url = url
+        self.model = model
+        self.temperature = temperature
+
+    def get_langchain_llm(self):
+        from langchain_ollama import ChatOllama
+
+        return ChatOllama(
+            base_url=self.url,
+            model=self.model,
+            temperature=self.temperature,
+        )
+
+
+@typechecked
+class LLMProvider(Enum):
+    OPENAI = "openai"
+    OLLAMA = "ollama"
+
+
+@typechecked
+class LLMManager(BaseModel):
+    """
+    Connection manager for Language Learning Models (LLMs).
+    """
+
+    url: str
+    provider: str
+    model: str
+    temperature: float = 0.0
+    api_key: Optional[str] = None
+
+    def connect(self) -> LLM:
+        """
+        Connect to the LLM.
+        """
+
+        match self.provider:
+            case LLMProvider.OPENAI.value:
+                if self.api_key is None:
+                    raise ValueError("API key is required for OpenAI LLM")
+                return OpenAILLM(
+                    api_key=self.api_key,
+                    model=self.model,
+                    temperature=self.temperature,
+                )
+            case LLMProvider.OLLAMA.value:
+                return OllamaLLM(
+                    url=self.url,
+                    model=self.model,
+                    temperature=self.temperature,
+                )
+            case _:
+                raise ValueError(f"Unknown LLM provider: {self.provider}")

uncertainty_engine_types-0.0.2/uncertainty_engine_types/message.py

@@ -0,0 +1,9 @@
+from datetime import datetime
+
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+    timestamp: datetime

uncertainty_engine_types-0.0.2/uncertainty_engine_types/model.py

@@ -0,0 +1,39 @@
+from typing import Any, Dict, Tuple
+import json
+import tempfile
+
+from pydantic import BaseModel, ConfigDict
+
+
+class TwinLabModel(BaseModel):
+    model_type: str
+    config: dict
+    metadata: dict
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    def load_model(self) -> Tuple[Any, Dict]:
+
+        from twinlab_models.models import model_type_from_str  # type: ignore
+
+        model_type = self.model_type
+        tl_model = model_type_from_str(model_type)
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".json") as f:
+            # Write the JSON to the temporary file
+            json.dump(self.model_dump(), f)  # Changed dumps to dump
+            f.flush()
+            tl_model, meta_data = tl_model.load(f.name)
+
+        return tl_model, meta_data
+
+
+# TODO: Should this be a method of the type?
+def save_model(model, meta_data: dict) -> TwinLabModel:
+
+    with tempfile.NamedTemporaryFile(mode="r", suffix=".json") as f:
+        model.save(f.name, meta_data)
+        f.seek(0)
+        config = json.load(f)
+
+    return TwinLabModel(**config)

uncertainty_engine_types-0.0.2/uncertainty_engine_types/node_info.py

@@ -0,0 +1,30 @@
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+
+class NodeInputInfo(BaseModel):
+    label: str
+    type: str
+    description: str
+    required: bool = True
+    set_in_node: bool = True
+    default: Optional[Any] = None
+
+
+class NodeOutputInfo(BaseModel):
+    label: str
+    type: str
+    description: str
+
+
+class NodeInfo(BaseModel):
+    id: str
+    version: int
+    label: str
+    category: str
+    description: str
+    long_description: str
+    cost: int
+    inputs: dict[str, NodeInputInfo]
+    outputs: dict[str, NodeOutputInfo] = {}

uncertainty_engine_types-0.0.2/uncertainty_engine_types/sensor_designer.py

@@ -0,0 +1,19 @@
+import json
+
+from pydantic import BaseModel
+
+
+class SensorDesigner(BaseModel):
+    bed: dict
+
+    def load_sensor_designer(self):
+
+        from twinlab_bed.BED import BED
+
+        bed = BED.from_json(json.dumps(self.bed))
+
+        return bed
+
+
+def save_sensor_designer(bed) -> SensorDesigner:
+    return SensorDesigner(bed=json.loads(bed.to_json()))

uncertainty_engine_types-0.0.2/uncertainty_engine_types/sql.py

@@ -0,0 +1,84 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel
+from typeguard import typechecked
+
+
+class SQLDatabase(ABC):
+    @abstractmethod
+    def execute(self, query: str) -> list[tuple[Any, ...]]:
+        """
+        Execute the given query.
+
+        Parameters:
+            query (str): The SQL query to execute.
+
+        Returns:
+            str: The result of the query.
+        """
+        pass
+
+
+class PostgreSQL(SQLDatabase):
+    @typechecked
+    def __init__(
+        self, host: str, username: str, password: str, port: int, database: str
+    ):
+        import psycopg2
+
+        self.connection = psycopg2.connect(
+            host=host, user=username, password=password, port=port, database=database
+        )
+        self.cursor = self.connection.cursor()
+
+    @typechecked
+    def execute(self, query: str) -> list[tuple[Any, ...]]:
+        """
+        Execute the given query.
+
+        Parameters:
+            query (str): The SQL query to execute.
+
+        Returns:
+            str: The result of the query.
+        """
+
+        self.cursor.execute(query)
+        return self.cursor.fetchall()
+
+
+class SQLKind(Enum):
+    POSTGRES = "POSTGRES"
+
+
+class SQLManager(BaseModel):
+    """
+    Connection manager for SQL databases.
+    """
+
+    kind: SQLKind
+    host: str
+    username: str
+    password: str
+    port: int
+    database: str
+
+    @typechecked
+    def connect(self) -> SQLDatabase:
+        """
+        Connect to the SQL database.
+        """
+
+        match self.kind:
+            case SQLKind.POSTGRES:
+                return PostgreSQL(
+                    host=self.host,
+                    username=self.username,
+                    password=self.password,
+                    port=self.port,
+                    database=self.database,
+                )
+            case _:
+                raise ValueError("Unsupported SQL kind")

uncertainty_engine_types-0.0.2/uncertainty_engine_types/tabular_data.py

@@ -0,0 +1,13 @@
+from io import StringIO
+
+from pydantic import BaseModel
+
+
+class TabularData(BaseModel):
+    csv: str
+
+    def load_dataframe(self):
+
+        import pandas as pd
+
+        return pd.read_csv(StringIO(self.csv))

uncertainty_engine_types-0.0.2/uncertainty_engine_types/token.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class Token(str, Enum):
+    TRAINING = "Training"  # Only used to mark nodes used for training models
+    STANDARD = "Standard"  # The default token type

uncertainty_engine_types-0.0.2/uncertainty_engine_types/vector_store.py

@@ -0,0 +1,230 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Optional, Literal, List, Dict
+
+from pydantic import BaseModel
+from typeguard import typechecked
+
+
+class VectorStoreProvider(Enum):
+    WEAVIATE = "weaviate"
+
+
+class VectorStoreConnection(ABC):
+    @abstractmethod
+    def ingest(self, texts, metadatas):
+        pass
+
+    @abstractmethod
+    def retrieve(self, query, k):
+        pass
+
+    @abstractmethod
+    def get_vector_store(self):
+        pass
+
+    @abstractmethod
+    def close(self):
+        pass
+
+
+@typechecked
+class WeaviateVectorStoreConnection(VectorStoreConnection):
+
+    def __init__(
+        self,
+        host: str,
+        port: str,
+        collection: str,
+        embedding_type: str,
+        embedding_model: str,
+        embedding_api_key: str,
+    ):
+
+        self.host = host
+        self.port = port
+        self.collection = collection
+        self.embedding_type = embedding_type
+        self.embedding_model = embedding_model
+        self.embedding_api_key = embedding_api_key
+        self.vector_store = get_persistent_vector_store(
+            host, port, collection, embedding_type, embedding_model, embedding_api_key
+        )
+
+    def ingest(self, texts: List[str], metadatas: List[Dict]) -> List[str]:
+        """
+        Ingest a list of texts into the vector store.
+
+        Args:
+            texts (List[str]): The texts to ingest.
+            metadatas (List[Dict]): The metadata associated with each text.
+
+        Returns:
+            List[str]: The IDs of the ingested texts.
+
+        """
+
+        return self.vector_store.add_texts(texts=texts, metadatas=metadatas)
+
+    def retrieve(self, query: Optional[str], k: int) -> List[Dict]:
+        """
+        Retrieve the k most relevant documents to a query.
+
+        Args:
+            query (Optional[str]): The query to retrieve documents for.
+            k (int): The number of documents to retrieve.
+
+        Returns:
+            List[str]: The IDs of the retrieved documents.
+
+        """
+
+        docs = self.vector_store.max_marginal_relevance_search(query=query, k=k)
+        docs_list = [
+            {"content": doc.page_content, "metadata": doc.metadata} for doc in docs
+        ]
+        return docs_list
+
+    def get_vector_store(self):
+        """
+        Get the underlying vector store.
+
+        Returns:
+            LangcahinWeaviateVectorStore: The underlying vector store.
+
+        """
+
+        return self.vector_store
+
+    def close(self):
+        self.vector_store._client.close()
+
+
+class VectorStoreManager(BaseModel):
+    """
+    Connection manager for a vector store.
+    """
+
+    provider: str
+    host: str
+    port: str = "8080"
+    collection: str = "DefaultCollection"
+    embedding_type: str
+    embedding_model: str
+    embedding_api_key: str
+
+    @typechecked
+    def connect(self) -> WeaviateVectorStoreConnection:
+        """
+        Connect to the vector store.
+
+        Returns:
+            VectorStoreConnection: The vector store connection.
+        """
+
+        match self.provider:
+            case VectorStoreProvider.WEAVIATE.value:
+                return WeaviateVectorStoreConnection(
+                    host=self.host,
+                    port=self.port,
+                    collection=self.collection,
+                    embedding_type=self.embedding_type,
+                    embedding_model=self.embedding_model,
+                    embedding_api_key=self.embedding_api_key,
+                )
+            case _:
+                raise ValueError(f"Unknown vector store provider: {self.provider}")
+
+
+@typechecked
+def get_persistent_vector_store(
+    host: str,
+    port: str,
+    collection: str,
+    embedding_type: str,
+    embedding_model: str,
+    embedding_api_key: Optional[str] = None,
+):
+    """
+    Get a database client connected to a deployed Weaviate vector store
+    """
+
+    from langchain_weaviate import WeaviateVectorStore
+    from weaviate.connect import ConnectionParams
+    import weaviate
+
+    try:
+        client = weaviate.WeaviateClient(
+            connection_params=ConnectionParams.from_params(
+                http_host=host,
+                http_port=port,
+                http_secure=False,
+                grpc_host=host,
+                grpc_port="50051",
+                grpc_secure=False,
+            ),
+            skip_init_checks=True,
+        )
+        client.connect()
+    except Exception as e:
+        raise ValueError(f"Failed to connect to Weaviate: {e}")
+
+    embedding_function = get_embedding_function(
+        embedding_type, embedding_model, embedding_api_key
+    )
+
+    # check if a collection exists and makes one if it doesn't
+    if not client.collections.exists(collection):
+        client.collections.create(collection)
+
+    try:
+        vector_store = WeaviateVectorStore(
+            client=client,
+            index_name=collection,
+            text_key="text",
+            embedding=embedding_function,
+        )
+    except Exception as e:
+        raise ValueError(f"Failed to initialize Weaviate vector store: {e}")
+
+    return vector_store
+
+
+@typechecked
+def get_embedding_function(
+    embedding_type: Literal["huggingface", "openai"],
+    embedding_model: str,
+    embedding_api_key: Optional[str] = None,
+):
+    """
+    Get an embedding function based on the specified type and configuration.
+    """
+
+    from langchain_huggingface import HuggingFaceEmbeddings
+    from langchain_openai import OpenAIEmbeddings
+
+    match embedding_type:
+        case "huggingface":
+            try:
+                return HuggingFaceEmbeddings(model_name=embedding_model)
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to initialize HuggingFace embedding model: {str(e)}"
+                ) from e
+
+        case "openai":
+            if not embedding_api_key:
+                raise ValueError("OpenAI embeddings require an API key")
+            try:
+                return OpenAIEmbeddings(
+                    model=embedding_model, api_key=embedding_api_key
+                )
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to initialize OpenAI embedding model: {str(e)}"
+                ) from e
+
+        case _:
+            raise ValueError(
+                f"Embedding type must be one of ['huggingface', 'openai']. Got {embedding_type}"
+            )