vectorwave 0.1.0__py3-none-any.whl

tests/models/test_db_config.py

@@ -0,0 +1,143 @@
+import pytest
+import json
+from unittest.mock import patch, mock_open
+from json import JSONDecodeError
+
+# Function to test
+from vectorwave.models.db_config import get_weaviate_settings
+
+# --- Mock Data ---
+
+# Mock content for a successfully loaded .weaviate_properties file
+MOCK_JSON_DATA = """
+{
+  "run_id": {
+    "data_type": "TEXT",
+    "description": "Test run ID"
+  },
+  "experiment_id": {
+    "data_type": "INT",
+    "description": "Identifier for the experiment"
+  }
+}
+"""
+
+# Mock content for a malformed .weaviate_properties file (invalid JSON)
+MOCK_INVALID_JSON = """
+{
+  "run_id": {
+    "data_type": "TEXT"
+  } 
+""" # Missing closing '}'
+
+# --- Test Cases ---
+
+@patch('os.path.exists', return_value=True)
+@patch('builtins.open', new_callable=mock_open, read_data=MOCK_JSON_DATA)
+def test_get_settings_loads_custom_props_success(mock_open_file, mock_exists):
+    """
+    Case 1: .weaviate_properties file exists and JSON is valid
+    - settings.custom_properties should be loaded correctly as a dictionary
+    """
+    # Arrange
+    # Clear the @lru_cache to bypass caching for this test
+    get_weaviate_settings.cache_clear()
+
+    # Act
+    settings = get_weaviate_settings()
+
+    # Assert
+    # Verify that the default path (.weaviate_properties) was checked
+    mock_exists.assert_called_with(".weaviate_properties")
+    # Verify the file was opened in 'r' mode
+    mock_open_file.assert_called_with(".weaviate_properties", 'r', encoding='utf-8')
+
+    assert settings.custom_properties is not None
+    assert "run_id" in settings.custom_properties
+    assert settings.custom_properties["run_id"]["data_type"] == "TEXT"
+    assert settings.custom_properties["run_id"]["description"] == "Test run ID"
+    assert "experiment_id" in settings.custom_properties
+
+
+@patch('os.path.exists', return_value=False)
+def test_get_settings_file_not_found(mock_exists, capsys):
+    """
+    Case 2: .weaviate_properties file does not exist
+    - settings.custom_properties should be None
+    - A 'file not found' note should be printed
+    """
+    # Arrange
+    get_weaviate_settings.cache_clear()
+
+    # Act
+    settings = get_weaviate_settings()
+
+    # Assert
+    mock_exists.assert_called_with(".weaviate_properties")
+    assert settings.custom_properties is None
+
+    # Check if 'file not found' note was printed
+    captured = capsys.readouterr()
+    assert "file not found" in captured.out
+
+
+@patch('os.path.exists', return_value=True)
+@patch('builtins.open', new_callable=mock_open, read_data=MOCK_INVALID_JSON)
+@patch('json.load', side_effect=JSONDecodeError("Mock JSON Decode Error", "", 0))
+def test_get_settings_invalid_json(mock_json_load, mock_open_file, mock_exists, capsys):
+    """
+    Case 3: File exists but JSON format is invalid
+    - settings.custom_properties should be None
+    - A 'Could not parse JSON' warning should be printed
+    """
+    # Arrange
+    get_weaviate_settings.cache_clear()
+
+    # Act
+    settings = get_weaviate_settings()
+
+    # Assert
+    mock_exists.assert_called_once()
+    mock_open_file.assert_called_once()
+    mock_json_load.assert_called_once() # json.load was called but failed (due to side_effect)
+    assert settings.custom_properties is None # Should be None due to parsing failure
+
+    # Check if 'Could not parse JSON' warning was printed
+    captured = capsys.readouterr()
+    assert "Could not parse JSON" in captured.out
+
+@patch('os.path.exists', return_value=True)
+@patch('builtins.open', new_callable=mock_open, read_data=MOCK_JSON_DATA)
+@patch('os.environ.get') # os.environ.get을 모킹합니다
+def test_get_settings_loads_global_custom_values(mock_env_get, mock_open_file, mock_exists):
+    """
+    Case 4: Test if the value of the "RUN_ID" environment variable is loaded
+    into global_custom_values for "run_id" defined in .weaviate_properties
+    """
+    # 1. Arrange
+    # MOCK_JSON_DATA defines "run_id" and "experiment_id".
+    # Set os.environ.get("RUN_ID") to return "test-run-123".
+    # Set os.environ.get("EXPERIMENT_ID") to return None.
+    def mock_env_side_effect(key):
+        if key == "RUN_ID":
+            return "test-run-123"
+        return None
+
+    mock_env_get.side_effect = mock_env_side_effect
+    get_weaviate_settings.cache_clear()
+
+    # 2. Act
+    settings = get_weaviate_settings()
+
+    # 3. Assert
+    # .weaviate_properties should be loaded correctly.
+    assert settings.custom_properties is not None
+    assert "run_id" in settings.custom_properties
+
+    # Check if global_custom_values was loaded correctly.
+    assert settings.global_custom_values is not None
+    assert "run_id" in settings.global_custom_values
+    assert settings.global_custom_values["run_id"] == "test-run-123"
+
+    # "EXPERIMENT_ID" should not be included as os.environ.get returned None.
+    assert "experiment_id" not in settings.global_custom_values

vectorwave/__init__.py

@@ -0,0 +1,13 @@
+from .core.decorator import vectorize
+
+from .database.db import initialize_database
+from .database.db_search import search_functions, search_executions
+from .monitoring.tracer import trace_span
+
+__all__ = [
+    'vectorize',
+    'initialize_database',
+    'search_functions',
+    'search_executions',
+    'trace_span'
+]

vectorwave/batch/batch.py

@@ -0,0 +1,65 @@
+import weaviate
+import atexit
+import logging
+from functools import lru_cache
+from ..models.db_config import get_weaviate_settings, WeaviateSettings
+from ..database.db import get_weaviate_client
+from ..exception.exceptions import WeaviateConnectionError
+
+
+
+class WeaviateBatchManager:
+    """
+    A singleton class that manages Weaviate batch imports.
+    """
+
+    def __init__(self):
+        self._initialized = False
+        print("Initializing WeaviateBatchManager...")
+        self.client: weaviate.WeaviateClient = None
+
+        try:
+            # (get_weaviate_settings is reused as it is handled by lru_cache)
+            self.settings: WeaviateSettings = get_weaviate_settings()
+            self.client: weaviate.WeaviateClient = get_weaviate_client(self.settings)
+
+            if not self.client:
+                raise WeaviateConnectionError("Client is None, cannot configure batch.")
+
+            # self.client.batch.configure(
+            #     batch_size=20,
+            #     dynamic=True,
+            #     timeout_retries=3,
+            # )
+
+            # Register atexit: Automatically calls self.flush() on script exit
+            # atexit.register(self.flush)
+            self._initialized = True
+            print("WeaviateBatchManager initialized and 'atexit' flush registered.")
+
+        except Exception as e:
+            # Prevents VectorWave from stopping the main app upon DB connection failure
+            print(f"Failed to initialize WeaviateBatchManager: {e}. Batching will be disabled.")
+
+    def add_object(self, collection: str, properties: dict, uuid: str = None):
+        """
+        Adds an object to the Weaviate batch queue.
+        """
+        if not self._initialized or not self.client:
+            print("Warning: Batch manager not initialized. Skipping add_object.") # print 유지
+            return
+
+        try:
+            self.client.collections.get(collection).data.insert(
+                properties=properties,
+                uuid=uuid
+            )
+
+        except Exception as e:
+            print(f"Error: Failed to add object to batch (Collection: {collection}): {e}")
+
+
+
+@lru_cache(None)
+def get_batch_manager() -> WeaviateBatchManager:
+    return WeaviateBatchManager()

vectorwave/core/decorator.py

@@ -0,0 +1,108 @@
+# vtm/src/vectorwave/core/decorator.py
+
+import inspect
+from functools import wraps
+
+from weaviate.util import generate_uuid5
+
+from ..batch.batch import get_batch_manager
+from ..models.db_config import get_weaviate_settings
+from ..monitoring.tracer import trace_root, trace_span
+
+
+def vectorize(search_description: str,
+              sequence_narrative: str,
+              **execution_tags):
+    """
+    VectorWave Decorator
+
+    (1) Collects function definitions (static data) once on script load.
+    (2) Records function execution (dynamic data) every time the function is called.
+    """
+
+    def decorator(func):
+
+        func_uuid = None
+        valid_execution_tags = {}
+        try:
+            module_name = func.__module__
+            function_name = func.__name__
+
+            func_identifier = f"{module_name}.{function_name}"
+            func_uuid = generate_uuid5(func_identifier)
+
+            static_properties = {
+                "function_name": function_name,
+                "module_name": module_name,
+                "docstring": inspect.getdoc(func) or "",
+                "source_code": inspect.getsource(func),
+                "search_description": search_description,
+                "sequence_narrative": sequence_narrative
+            }
+
+            batch = get_batch_manager()
+            settings = get_weaviate_settings()
+
+            if execution_tags:
+                if not settings.custom_properties:
+                    print(
+                        f"Warning: Function '{function_name}' provided execution_tags {list(execution_tags.keys())} "
+                        f"but no .weaviate_properties file was loaded. These tags will be IGNORED."
+                    )
+                else:
+                    allowed_keys = set(settings.custom_properties.keys())
+                    for key, value in execution_tags.items():
+                        if key in allowed_keys:
+                            valid_execution_tags[key] = value
+                        else:
+                            print(
+                                f"Warning: Function '{function_name}' has undefined execution_tag: '{key}'. "
+                                f"This tag will be IGNORED. Please add it to your .weaviate_properties file."
+                            )
+
+            batch.add_object(
+                collection=settings.COLLECTION_NAME,
+                properties=static_properties,
+                uuid=func_uuid
+            )
+
+        except Exception as e:
+            print(f"Error in @vectorize setup for {func.__name__}: {e}")
+
+
+        # 2a. The *inner* wrapper to be wrapped by @trace_span
+        # This function receives all tags including full_kwargs from @trace_span.
+        @trace_root()
+        @trace_span(attributes_to_capture=['function_uuid', 'team', 'priority', 'run_id'])
+        @wraps(func)
+        def inner_wrapper(*args, **kwargs):
+
+            original_kwargs = kwargs.copy()
+
+            keys_to_remove = list(valid_execution_tags.keys())
+            keys_to_remove.append('function_uuid')
+
+            for key in execution_tags.keys():
+                if key not in keys_to_remove:
+                    keys_to_remove.append(key)
+
+            for key in keys_to_remove:
+                original_kwargs.pop(key, None)
+
+            return func(*args, **original_kwargs)
+
+
+        @wraps(func)
+        def outer_wrapper(*args, **kwargs):
+
+            full_kwargs = kwargs.copy()
+            full_kwargs.update(valid_execution_tags)
+            full_kwargs['function_uuid'] = func_uuid
+
+            # 2. Call the *inner* wrapper with the full_kwargs
+            #    This call passes through the @trace_root -> @trace_span decorators.
+            return inner_wrapper(*args, **full_kwargs)
+
+        return outer_wrapper
+
+    return decorator

vectorwave/database/db.py

@@ -0,0 +1,302 @@
+import weaviate
+import weaviate.classes.config as wvc # (wvc = Weaviate Classes Config)
+import weaviate.config as wvc_config
+from weaviate.config import AdditionalConfig
+from vectorwave.models.db_config import WeaviateSettings
+from vectorwave.exception.exceptions import (
+    WeaviateConnectionError,
+    WeaviateNotReadyError,
+    SchemaCreationError
+)
+from functools import lru_cache
+from weaviate.exceptions import WeaviateConnectionError as WeaviateClientConnectionError
+from vectorwave.models.db_config import get_weaviate_settings
+
+
+# Code based on Weaviate v4 (latest) client.
+
+def get_weaviate_client(settings: WeaviateSettings) -> weaviate.WeaviateClient:
+    """
+    Creates and returns a Weaviate client.
+
+    [Raises]
+    - WeaviateConnectionError: If connection fails.
+    - WeaviateNotReadyError: If connected, but the server is not ready.
+    """
+
+    client: weaviate.WeaviateClient
+
+    try:
+        client = weaviate.connect_to_local(
+            host=settings.WEAVIATE_HOST,
+            port=settings.WEAVIATE_PORT,
+            grpc_port=settings.WEAVIATE_GRPC_PORT,
+            additional_config=AdditionalConfig(
+                    dynamic=True,
+                    batch_size=20,
+                    timeout_retries=3
+            )
+        )
+    except WeaviateClientConnectionError as e:
+        raise WeaviateConnectionError(f"Failed to connect to Weaviate: {e}")
+    except Exception as e:
+        raise WeaviateConnectionError(f"An unknown error occurred while connecting to Weaviate: {e}")
+
+
+    if not client.is_ready():
+        raise WeaviateNotReadyError("Connected to Weaviate, but the server is not ready.")
+
+    print("Weaviate client connected successfully.")
+    return client
+
+@lru_cache()
+def get_cached_client() -> weaviate.WeaviateClient:
+    """
+    Singleton factory: Gets settings and returns a single client instance.
+    This function IS cached.
+    """
+    print("Creating and caching new Weaviate client instance...")
+    settings = get_weaviate_settings()
+    client = get_weaviate_client(settings)
+    return client
+
+
+def create_vectorwave_schema(client: weaviate.WeaviateClient, settings: WeaviateSettings):
+    """
+    Defines and creates the VectorWaveFunctions collection schema.
+    Now includes custom properties loaded from the settings (via .weaviate_properties file).
+
+    [Raises]
+    - SchemaCreationError: If an error occurs during schema creation.
+    """
+    collection_name = settings.COLLECTION_NAME
+
+    # 1. Check if the collection already exists
+    if client.collections.exists(collection_name):
+        print(f"Collection '{collection_name}' already exists. Skipping creation.")
+        return client.collections.get(collection_name)
+
+    # 2. If it doesn't exist, define and create the collection
+    print(f"Creating collection '{collection_name}'...")
+
+    # 3. Define Base Properties
+    base_properties = [
+        wvc.Property(
+            name="function_name",
+            data_type=wvc.DataType.TEXT,
+            description="The name of the vectorized function"
+        ),
+        wvc.Property(
+            name="module_name",
+            data_type=wvc.DataType.TEXT,
+            description="The Python module path where the function is defined"
+        ),
+        wvc.Property(
+            name="docstring",
+            data_type=wvc.DataType.TEXT,
+            description="The function's Docstring (description)"
+        ),
+        wvc.Property(
+            name="source_code",
+            data_type=wvc.DataType.TEXT,
+            description="The actual source code of the function"
+        ),
+        wvc.Property(
+            name="search_description",
+            data_type=wvc.DataType.TEXT,
+            description="User-provided description for similarity search (from @vectorize)"
+        ),
+        wvc.Property(
+            name="sequence_narrative",
+            data_type=wvc.DataType.TEXT,
+            description="User-provided context about what happens next (from @vectorize)"
+        ),
+    ]
+
+    # 4. Parse Custom Properties (loaded from JSON file via settings object)
+    custom_properties = []
+    if settings.custom_properties:
+        print(f"Adding custom properties: {list(settings.custom_properties.keys())}")
+
+        for name, prop_details in settings.custom_properties.items():
+            if not isinstance(prop_details, dict):
+                raise SchemaCreationError(f"Custom property '{name}' in config file must be a dictionary.")
+
+            # Get data_type (Required)
+            dtype_str = prop_details.get("data_type")
+            if not dtype_str:
+                raise SchemaCreationError(f"Custom property '{name}' in config file is missing 'data_type'.")
+
+            # Get description (Optional)
+            description = prop_details.get("description")
+
+            try:
+                # Convert string (e.g., "TEXT") to Weaviate Enum (wvc.DataType.TEXT)
+                data_type = getattr(wvc.DataType, dtype_str.upper())
+
+                custom_properties.append(
+                    wvc.Property(
+                        name=name,
+                        data_type=data_type,
+                        description=description
+                    )
+                )
+            except AttributeError:
+                raise SchemaCreationError(
+                    f"Invalid data_type '{dtype_str}' for custom property '{name}'. "
+                    f"Use a valid wvc.DataType string (e.g., 'TEXT', 'INT', 'NUMBER')."
+                )
+            except Exception as e:
+                raise SchemaCreationError(f"Error processing custom property '{name}': {e}")
+
+    # 5. Combine properties
+    all_properties = base_properties + custom_properties
+
+    vectorizer_name = settings.VECTORIZER_CONFIG.lower()
+    vector_config = None
+
+    print(f"Configuring vectorizer: {vectorizer_name}")
+
+    if vectorizer_name == "text2vec-openai":
+        vector_config = {
+            "vectorizer": "text2vec-openai",
+            "text2vec-openai": {
+                "vectorizeClassName": settings.IS_VECTORIZE_COLLECTION_NAME,
+            }
+        }
+    elif vectorizer_name == "none":
+        vector_config = {"vectorizer": "none"}
+    else:
+        raise SchemaCreationError(
+            f"Unsupported VECTORIZER_CONFIG: '{settings.VECTORIZER_CONFIG}'. "
+            f"Supported values: 'text2vec-openai', 'none'."
+        )
+
+    generative_config = None
+    if settings.GENERATIVE_CONFIG.lower() == "generative-openai":
+        generative_config = {"generator": "generative-openai"}
+
+
+    try:
+        vectorwave_collection = client.collections.create(
+            name=collection_name,
+            properties=all_properties,
+
+            # 7. Vectorizer Configuration
+            vector_config=vector_config,
+
+            # 8. Generative Configuration (for RAG, etc.)
+            generative_config=generative_config
+        )
+        return vectorwave_collection
+
+    except Exception as e:
+        # Raise a specific exception instead of returning None
+        raise SchemaCreationError(f"Error during schema creation: {e}")
+
+
+def create_execution_schema(client: weaviate.WeaviateClient, settings: WeaviateSettings):
+    """
+    Defines and creates the VectorWaveExecutions (dynamic) collection schema.
+    """
+    collection_name = settings.EXECUTION_COLLECTION_NAME
+
+    if client.collections.exists(collection_name):
+        print(f"Collection '{collection_name}' already exists. Skipping creation.")
+        return client.collections.get(collection_name)
+
+    print(f"Creating collection '{collection_name}'...")
+
+    properties = [
+        wvc.Property(
+            name="trace_id",
+            data_type=wvc.DataType.TEXT,
+            description="The unique ID for the entire trace/workflow"
+        ),
+        wvc.Property(
+            name="span_id",
+            data_type=wvc.DataType.TEXT,
+            description="The unique ID for this specific span/function execution"
+        ),
+        wvc.Property(
+            name="function_name",
+            data_type=wvc.DataType.TEXT,
+            description="Name of the executed function (span name)"
+        ),
+        wvc.Property(
+            name="function_uuid",
+            data_type=wvc.DataType.UUID,
+            description="The UUID of the executed function definition"
+        ),
+        wvc.Property(
+            name="timestamp_utc",
+            data_type=wvc.DataType.DATE,
+            description="The UTC timestamp when the execution started"
+        ),
+        wvc.Property(
+            name="duration_ms",
+            data_type=wvc.DataType.NUMBER,
+            description="Total execution time in milliseconds"
+        ),
+        wvc.Property(
+            name="status",
+            data_type=wvc.DataType.TEXT, # "SUCCESS" or "ERROR"
+            description="Execution status"
+        ),
+        wvc.Property(
+            name="error_message",
+            data_type=wvc.DataType.TEXT,
+            description="Error message and traceback if status is 'ERROR'"
+        ),
+    ]
+
+    if settings.custom_properties:
+        print(f"Adding custom properties to '{collection_name}': {list(settings.custom_properties.keys())}")
+        for name, prop_details in settings.custom_properties.items():
+            try:
+                if not isinstance(prop_details, dict):
+                    raise ValueError("Property details must be a dictionary.")
+
+                dtype_str = prop_details.get("data_type")
+                if not dtype_str:
+                    raise ValueError("data_type is missing.")
+
+                data_type = getattr(wvc.DataType, dtype_str.upper())
+                description = prop_details.get("description")
+
+                properties.append(
+                    wvc.Property(
+                        name=name,
+                        data_type=data_type,
+                        description=description
+                    )
+                )
+            except Exception as e:
+                print(f"Warning: Skipping custom property {name} for {collection_name}: {e}")
+
+    try:
+        execution_collection = client.collections.create(
+            name=collection_name,
+            properties=properties,
+            vectorizer_config=wvc.Configure.Vectorizer.none(),
+            vector_index_config=wvc.Configure.VectorIndex.none()
+        )
+        print(f"Collection '{collection_name}' created successfully.")
+        return execution_collection
+    except Exception as e:
+        raise SchemaCreationError(f"Error during execution schema creation: {e}")
+
+def initialize_database():
+    """
+    Helper function to initialize both the client and the two schemas.
+    """
+    try:
+        settings = get_weaviate_settings()
+        client = get_cached_client()
+        if client:
+            create_vectorwave_schema(client, settings)
+            create_execution_schema(client, settings)
+            return client
+    except Exception as e:
+        print(f"Failed to initialize VectorWave database: {e}")
+        return None