crewplus 0.2.39__tar.gz → 0.2.41__tar.gz

{crewplus-0.2.39 → crewplus-0.2.41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.39
+Version: 0.2.41
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT

crewplus-0.2.41/crewplus/callbacks/__init__.py

@@ -0,0 +1 @@
+# This file makes the 'callbacks' directory a Python package.

crewplus-0.2.41/crewplus/callbacks/async_langfuse_handler.py

@@ -0,0 +1,95 @@
+# File: crewplus/callbacks/async_langfuse_handler.py
+import asyncio
+import contextvars
+from contextlib import contextmanager
+from typing import Any, Dict, List, Union
+
+try:
+    from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
+    from langchain_core.callbacks import AsyncCallbackHandler
+    from langchain_core.outputs import LLMResult
+    LANGFUSE_AVAILABLE = True
+except ImportError:
+    LANGFUSE_AVAILABLE = False
+    LangfuseCallbackHandler = None
+    AsyncCallbackHandler = object
+
+# This token is a simple flag to indicate that we are in an async context.
+# We use a context variable to make it available only within the async task.
+_ASYNC_CONTEXT_TOKEN = "in_async_context"
+in_async_context = contextvars.ContextVar(_ASYNC_CONTEXT_TOKEN, default=False)
+
+@contextmanager
+def async_context():
+    """A context manager to signal that we are in an async execution context."""
+    token = in_async_context.set(True)
+    try:
+        yield
+    finally:
+        in_async_context.reset(token)
+
+class AsyncLangfuseCallbackHandler(AsyncCallbackHandler):
+    """
+    Wraps the synchronous LangfuseCallbackHandler to make it compatible with
+    LangChain's async methods.
+    
+    This works by running the synchronous handler's methods in a separate thread
+    using `asyncio.to_thread`. This is crucial because `asyncio`'s default
+    executor can correctly propagate `contextvars`, which solves the
+    `ValueError: <Token ...> was created in a different Context` from OpenTelemetry.
+    """
+    def __init__(self, *args: Any, **kwargs: Any):
+        if not LANGFUSE_AVAILABLE:
+            raise ImportError("Langfuse is not available. Please install it with 'pip install langfuse'")
+        self.sync_handler = LangfuseCallbackHandler(*args, **kwargs)
+
+    def __getattr__(self, name: str) -> Any:
+        # Delegate any other attribute access to the sync handler
+        return getattr(self.sync_handler, name)
+
+    async def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        await asyncio.to_thread(
+            self.sync_handler.on_llm_start, serialized, prompts, **kwargs
+        )
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        await asyncio.to_thread(
+            self.sync_handler.on_llm_end, response, **kwargs
+        )
+    
+    async def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
+        await asyncio.to_thread(
+            self.sync_handler.on_llm_error, error, **kwargs
+        )
+    
+    async def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_tool_start, serialized, input_str, **kwargs
+        )
+
+    async def on_tool_end(self, output: str, **kwargs: Any) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_tool_end, output, **kwargs
+        )
+
+    async def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_tool_error, error, **kwargs
+        )
+
+    async def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_chain_start, serialized, inputs, **kwargs
+        )
+
+    async def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_chain_end, outputs, **kwargs
+        )
+
+    async def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_chain_error, error, **kwargs
+        )

{crewplus-0.2.39 → crewplus-0.2.41}/crewplus/services/azure_chat_model.py

@@ -5,6 +5,7 @@ from typing import Any, Optional
 from langchain_openai.chat_models.azure import AzureChatOpenAI
 from pydantic import Field
 from .tracing_manager import TracingManager, TracingContext
+from ..callbacks.async_langfuse_handler import async_context
 
 class TracedAzureChatOpenAI(AzureChatOpenAI):
     """
@@ -106,8 +107,9 @@ class TracedAzureChatOpenAI(AzureChatOpenAI):
         return super().invoke(input, config=config, **kwargs)
 
     async def ainvoke(self, input, config=None, **kwargs):
-        config = self._tracing_manager.add_callbacks_to_config(config)
-        return await super().ainvoke(input, config=config, **kwargs)
+        with async_context():
+            config = self._tracing_manager.add_callbacks_to_config(config)
+            return await super().ainvoke(input, config=config, **kwargs)
 
     def stream(self, input, config=None, **kwargs):
         # Add stream_options to get usage data for Langfuse
@@ -124,6 +126,7 @@ class TracedAzureChatOpenAI(AzureChatOpenAI):
         stream_options["include_usage"] = True
         kwargs["stream_options"] = stream_options
         
-        config = self._tracing_manager.add_callbacks_to_config(config)
-        async for chunk in super().astream(input, config=config, **kwargs):
-            yield chunk
+        with async_context():
+            config = self._tracing_manager.add_callbacks_to_config(config)
+            async for chunk in super().astream(input, config=config, **kwargs):
+                yield chunk

{crewplus-0.2.39 → crewplus-0.2.41}/crewplus/services/tracing_manager.py

@@ -8,10 +8,13 @@ import logging
 # even if the langfuse library is not installed.
 try:
     from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
+    from ..callbacks.async_langfuse_handler import AsyncLangfuseCallbackHandler, in_async_context
     LANGFUSE_AVAILABLE = True
 except ImportError:
     LANGFUSE_AVAILABLE = False
     LangfuseCallbackHandler = None
+    AsyncLangfuseCallbackHandler = None
+    in_async_context = None
 
 class TracingContext(Protocol):
     """
@@ -65,7 +68,8 @@ class TracingManager:
                      to the TracingContext protocol.
         """
         self.context = context
-        self._handlers: List[Any] = []
+        self._sync_handlers: List[Any] = []
+        self._async_handlers: List[Any] = []
         self._initialize_handlers()
     
     def _initialize_handlers(self):
@@ -73,7 +77,8 @@ class TracingManager:
         Initializes all supported tracing handlers. This is the central point
         for adding new observability tools.
         """
-        self._handlers = []
+        self._sync_handlers = []
+        self._async_handlers = []
         self._initialize_langfuse()
         # To add a new handler (e.g., Helicone), you would add a call to
         # self._initialize_helicone() here.
@@ -94,8 +99,14 @@ class TracingManager:
         
         if enable_langfuse:
             try:
-                handler = LangfuseCallbackHandler()
-                self._handlers.append(handler)
+                # Create both sync and async handlers. We'll pick one at runtime.
+                sync_handler = LangfuseCallbackHandler()
+                self._sync_handlers.append(sync_handler)
+
+                if AsyncLangfuseCallbackHandler:
+                    async_handler = AsyncLangfuseCallbackHandler()
+                    self._async_handlers.append(async_handler)
+                
                 self.context.logger.info(f"Langfuse tracing enabled for {self.context.get_model_identifier()}")
             except Exception as e:
                 self.context.logger.warning(f"Failed to initialize Langfuse: {e}")
@@ -118,15 +129,19 @@ class TracingManager:
         if config is None:
             config = {}
         
+        # Decide which handlers to use based on the async context flag.
+        is_async = in_async_context.get() if in_async_context else False
+        handlers = self._async_handlers if is_async else self._sync_handlers
+
         # Respect a global disable flag for this specific call.
-        if not self._handlers or config.get("metadata", {}).get("tracing_disabled"):
+        if not handlers or config.get("metadata", {}).get("tracing_disabled"):
             return config
         
         callbacks = config.get("callbacks")
         
         # Case 1: The 'callbacks' key holds a CallbackManager instance
         if hasattr(callbacks, 'add_handler') and hasattr(callbacks, 'handlers'):
-            for handler in self._handlers:
+            for handler in handlers:
                 if not any(isinstance(cb, type(handler)) for cb in callbacks.handlers):
                     callbacks.add_handler(handler, inherit=True)
             return config # Return the original, now-mutated config
@@ -135,7 +150,7 @@ class TracingManager:
         current_callbacks = callbacks or []
         new_callbacks = list(current_callbacks)
         
-        for handler in self._handlers:
+        for handler in handlers:
             if not any(isinstance(cb, type(handler)) for cb in new_callbacks):
                 new_callbacks.append(handler)
         

crewplus-0.2.41/crewplus/vectorstores/milvus/schema_milvus.py

@@ -0,0 +1,563 @@
+from typing import List, Optional
+import logging
+import json
+import asyncio
+
+from pymilvus import DataType
+from langchain_milvus import Milvus
+from langchain_core.documents import Document
+from ...utils.schema_document_updater import SchemaDocumentUpdater
+from ...utils.schema_action import Action
+from .milvus_schema_manager import MilvusSchemaManager
+
+DEFAULT_SCHEMA = """
+{
+    "node_types": {
+        "Document": {
+            "properties": {
+                "pk": {
+                    "type": "INT64",
+                    "is_primary": true,
+                    "auto_id": true
+                },
+                "vector": {
+                    "type": "FLOAT_VECTOR",
+                    "dim": 1536
+                },
+                "text": {
+                    "type": "VARCHAR",
+                    "max_length": 65535,
+                    "description": "The core text of the memory. This could be a user query, a documented fact, a procedural step, or a log of an event."
+                }
+	        }
+	    }
+    }
+}
+"""
+
+class SchemaMilvus(Milvus):
+    """
+    SchemaMilvus is a subclass of the Milvus class from langchain_milvus. This class is responsible for updating metadata of documents in a Milvus vector store.
+
+    Attributes:
+        embedding_function: Embedding function used by the Milvus vector store.
+        collection_name: Name of the collection in the Milvus vector store.
+        connection_args: Connection arguments for the Milvus vector store.
+        index_params: Index parameters for the Milvus vector store.
+        auto_id: Flag to specify if auto ID generation is enabled.
+        primary_field: The primary field of the collection.
+        vector_field: The vector field of the collection.
+        consistency_level: The consistency level for the Milvus vector store.
+        collection_schema: Schema JSON string associated with the Milvus existing collection name.
+    """
+    def __init__(
+        self, 
+        embedding_function, 
+        collection_name, 
+        connection_args, 
+        index_params=None, 
+        auto_id=True, 
+        primary_field="pk", 
+        text_field: str = "text",
+        vector_field=["vector"], 
+        consistency_level="Session",
+        logger: Optional[logging.Logger] = None
+    ):
+        """
+        Initializes the SchemaMilvus class with the provided parameters.
+
+        Args:
+            embedding_function: Embedding function used by the Milvus vector store.
+            collection_name: Name of the collection in the Milvus vector store.
+            connection_args: Connection arguments for the Milvus vector store.
+            index_params: Index parameters for the Milvus vector store.
+            auto_id: Flag to specify if auto ID generation is enabled.
+            primary_field: The primary field of the collection.
+            text_field: The text field of the collection.
+            vector_field: The vector field of the collection.
+            consistency_level: The consistency level for the Milvus vector store.
+            logger: Optional logger instance. If not provided, a default logger is created.
+        """
+        super().__init__(
+            embedding_function=embedding_function,
+            collection_name=collection_name,
+            connection_args=connection_args,
+            index_params=index_params,
+            auto_id=auto_id,
+            primary_field=primary_field,
+            text_field=text_field,
+            vector_field=vector_field,
+            consistency_level=consistency_level
+        )
+        self.logger = logger or logging.getLogger(__name__)
+        self.collection_schema = None
+        self.schema_manager = MilvusSchemaManager(client=self.client)
+
+    def set_schema(self, schema: str):
+        """
+        Sets the collection schema.
+
+        Args:
+            schema: The schema JSON string.
+        """
+        self.collection_schema = schema
+    
+    def get_fields(self, collection_name: Optional[str] = None) -> Optional[List[str]]:
+        """
+        Retrieves and returns the fields from the collection schema.
+
+        Args:
+            collection_name: The name of the collection to describe. If None, use self.collection_name.
+
+        Returns:
+            List[str] | None: The list of field names from the collection schema (excluding vector and text fields), or None if collection_name is not provided or an error occurs.
+        """
+        if collection_name is None:
+            collection_name = self.collection_name
+        if collection_name is None:
+            return None
+
+        try:
+            schema = self.client.describe_collection(collection_name)
+            fields = [field["name"] for field in schema["fields"] if field["type"] != DataType.FLOAT_VECTOR ]
+            return fields
+        except Exception as e:
+            self.logger.warning(f"Failed to retrieve schema fields: {e}")
+            return None
+    
+    def create_collection(self) -> bool:
+        """
+        Validates the schema and creates the collection using the MilvusSchemaManager.
+
+        Returns:
+            bool: True if the collection is successfully created, False otherwise.
+        """
+        if self.collection_schema is None:
+            self.logger.error("Collection schema is not set. Please set a schema using set_schema().")
+            return False
+            
+        self.schema_manager.bind_client(self.client)
+        if not self.schema_manager.validate_schema(self.collection_schema):
+            self.logger.error("Failed to validate schema")
+            return False
+        try:
+            self.schema_manager.create_collection(self.collection_name, self.collection_schema)
+            self.logger.info(f"Collection {self.collection_name} created successfully")
+            
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to create collection: {e}")
+            return False
+
+    def drop_collection(self, collection_name: Optional[str] = None) -> bool:
+        """
+        Drops the collection using the Milvus client.
+
+        Returns:
+            bool: True if the collection is successfully dropped, False otherwise.
+        """
+        if collection_name is None:
+            collection_name = self.collection_name
+            
+        try:
+            self.client.drop_collection(collection_name)
+            self.logger.info(f"Collection {collection_name} dropped successfully")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to drop collection {self.collection_name}: {e}")
+            return False
+
+    def _handle_upsert(self, doc: Document, metadata_dict: dict) -> Document:
+        """
+        Handles the UPSERT action for a single document by merging metadata.
+        """
+        existing_metadata = doc.metadata
+        for key, value in metadata_dict.items():
+            # Skip primary key and text fields to prevent modification.
+            if key in [self.primary_field, self.text_field]:
+                continue
+
+            if isinstance(value, dict):
+                # If it's a JSON object field (e.g., plant_metadata)
+                # Check if the existing value is a string, and if so, try to parse it as a dictionary
+                if key in existing_metadata and isinstance(existing_metadata[key], str):
+                    try:
+                        existing_metadata[key] = json.loads(existing_metadata[key])
+                    except json.JSONDecodeError:
+                        # If the parsing fails, it may not be a valid JSON string, treat it as a regular string
+                        self.logger.warning(f"Field '{key}' could not be parsed as JSON. Overwriting as a new dict.")
+                        existing_metadata[key] = {}
+
+                if key not in existing_metadata:
+                    # If the field does not exist, add it
+                    existing_metadata[key] = value
+                elif isinstance(existing_metadata[key], dict):
+                    # If the field exists and is a dictionary, recursively update the sub-fields
+                    for sub_key, sub_value in value.items():
+                        if isinstance(sub_value, dict):
+                            # If the sub-field is also a dictionary, recursively process it
+                            if sub_key not in existing_metadata[key]:
+                                existing_metadata[key][sub_key] = sub_value
+                            else:
+                                existing_metadata[key][sub_key].update(sub_value)
+                        else:
+                            # If the sub-field is a regular value, update it
+                            existing_metadata[key][sub_key] = sub_value
+                else:
+                    # If the field exists but is not a dictionary (e.g., a number or string), overwrite with the new dictionary
+                    existing_metadata[key] = value
+            else:
+                # If it's a regular field, update the value
+                existing_metadata[key] = value
+
+        # Update the document's metadata
+        doc.metadata = existing_metadata
+
+        return doc
+
+    def _process_document_update(self, doc: Document, metadata_dict: dict, action: Action) -> Document:
+        """
+        Applies the specified update operation to a single document.
+        
+        Args:
+            doc: The Document object to be updated.
+            metadata_dict: A dictionary containing the new data.
+            action: The type of operation to perform (UPSERT, DELETE, UPDATE, INSERT).
+
+        Returns:
+            The updated Document object.
+        """
+        pk_value = doc.metadata.get(self.primary_field)
+        text_value = doc.metadata.get(self.text_field)
+
+        if action == Action.UPSERT:
+            doc = self._handle_upsert(doc, metadata_dict)
+        elif action == Action.DELETE:
+            keys_to_delete = metadata_dict.keys()
+            doc = SchemaDocumentUpdater.delete_document_metadata(doc, list(keys_to_delete))
+        elif action == Action.UPDATE:
+            existing_metadata = doc.metadata
+            update_dict = {}
+            for key, value in metadata_dict.items():
+                if key in existing_metadata:
+                    if isinstance(value, dict) and isinstance(existing_metadata.get(key), dict):
+                        merged = existing_metadata[key].copy()
+                        for sub_key, sub_value in value.items():
+                            if sub_key in merged:
+                                merged[sub_key] = sub_value
+                        update_dict[key] = merged
+                    else:
+                        update_dict[key] = value
+            doc = SchemaDocumentUpdater.update_document_metadata(doc, update_dict)
+        elif action == Action.INSERT:
+            existing_metadata = doc.metadata
+            for key, value in metadata_dict.items():
+                if key in ['pk', 'text']:
+                    continue
+
+                if isinstance(value, dict) and key in existing_metadata and isinstance(existing_metadata.get(key), dict):
+                    existing_metadata[key] = {}
+                    existing_metadata[key] = value
+                else:
+                    existing_metadata[key] = value
+            doc.metadata = existing_metadata
+
+        if pk_value is not None:
+            doc.metadata[self.primary_field] = pk_value
+        if text_value is not None:
+            doc.metadata[self.text_field] = text_value
+        
+        return doc
+
+    def update_documents_metadata(self, expr: str, metadata: str, action: Action = Action.UPSERT) -> List[Document]:
+        """
+        Updates the metadata of documents in the Milvus vector store based on the provided expression.
+        This method uses a direct client upsert to avoid re-embedding vectors.
+
+        Args:
+            expr: Expression to filter the target documents.
+            metadata: New metadata to update the documents with.
+            action: The action to perform on the document metadata.
+
+        Returns:
+            List of updated documents.
+        """
+        try:
+            metadata_dict = json.loads(metadata)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON string for metadata")
+        
+        fields = self.get_fields()
+        if not fields:
+            fields = []
+
+        if isinstance(self._vector_field, list):
+            fields.extend(self._vector_field)
+        else:
+            fields.append(self._vector_field)
+
+        documents = self.search_by_metadata(expr, fields=fields, limit=5000)
+        
+        updated_documents = [self._process_document_update(doc, metadata_dict, action) for doc in documents]
+
+        if updated_documents:
+            self.logger.debug(f"Upserting {len(updated_documents)} documents using direct client upsert.")
+            upsert_data = [doc.metadata for doc in updated_documents]
+            self.client.upsert(
+                collection_name=self.collection_name,
+                data=upsert_data
+            )
+        
+        return updated_documents
+
+    async def aupdate_documents_metadata(self, expr: str, metadata: str, action: Action = Action.UPSERT) -> List[Document]:
+        """
+        Asynchronously updates the metadata of documents in the Milvus vector store.
+        This method uses a direct client upsert to avoid re-embedding vectors.
+
+        Args:
+            expr: Expression to filter the target documents.
+            metadata: New metadata to update the documents with.
+            action: The action to perform on the document metadata.
+
+        Returns:
+            List of updated documents.
+        """
+        try:
+            metadata_dict = json.loads(metadata)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON string for metadata")
+        
+        fields = self.get_fields()
+        if not fields:
+            fields = []
+
+        if isinstance(self._vector_field, list):
+            fields.extend(self._vector_field)
+        else:
+            fields.append(self._vector_field)
+
+        documents = self.search_by_metadata(expr, fields=fields, limit=5000)
+        
+        updated_documents = [self._process_document_update(doc, metadata_dict, action) for doc in documents]
+
+        if updated_documents:
+            self.logger.debug(f"Upserting {len(updated_documents)} documents using direct client upsert.")
+            upsert_data = [doc.metadata for doc in updated_documents]
+            
+            await asyncio.to_thread(
+                self.client.upsert,
+                collection_name=self.collection_name,
+                data=upsert_data
+            )
+        
+        return updated_documents
+
+    def update_documents_metadata_by_iterator(self, expr: str, metadata: str, action:Action=Action.UPSERT) -> List[Document]:
+        """
+        【官方推荐版】
+        使用 pymilvus.Collection.query_iterator 官方推荐的迭代方式更新元数据。
+        本方法的业务逻辑（UPSERT/DELETE等）与 update_documents_metadata 方法完全一致，
+        仅数据获取方式遵循官方标准迭代器模式。
+        """
+        try:
+            metadata_dict = json.loads(metadata)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON string for metadata")
+
+        fields = self.get_fields() or []
+        # 确保主键和文本字段在输出字段中
+        if 'pk' not in fields:
+            fields.append('pk')
+        text_field = getattr(self, "_text_field", "text")
+        if text_field not in fields:
+            fields.append(text_field)
+        
+        # 【关键修正】: 确保在查询时也获取向量字段。
+        # self.client.upsert 操作要求提供所有非 nullable 字段，包括 vector。
+        # 因此，我们必须在迭代查询时获取它，以便在更新时能够一并传回。
+        vector_fields_to_add = self._vector_field if isinstance(self._vector_field, list) else [self._vector_field]
+        for vf in vector_fields_to_add:
+            if vf not in fields:
+                fields.append(vf)
+            
+        total_updated_documents = []
+        batch_size = 1000 # 您可以根据需要调整批次大小
+
+        self.logger.info(f"Starting metadata update using 'collection.query_iterator' with batch size {batch_size}.")
+
+        # # 1. 【关键保险】: 在查询前，显式地确保集合已被加载。
+        # logger.info(f"Ensuring collection '{self.collection_name}' is loaded before querying.")
+        # self.col.load()
+
+        # 2. 【官方用法】: 创建官方推荐的迭代器
+        iterator = self.col.query_iterator(
+            batch_size=batch_size,
+            expr=expr,
+            output_fields=fields
+        )
+
+        batch_i = 0
+        try:
+            while True:
+                # 3. 【官方用法】: 获取下一批次
+                batch_results = iterator.next()
+                if not batch_results:
+                    break # 迭代完成，正常退出
+
+                batch_i += 1
+                self.logger.info(f"Processing batch {batch_i} of {len(batch_results)} documents.")
+
+                # 4. 将 Milvus 返回的 dict 列表转换为 Langchain Document 对象
+                documents = [
+                    Document(page_content=result.get(text_field, ""), metadata=result)
+                    for result in batch_results
+                ]
+
+                # 5. 【核心业务逻辑】: 使用公共方法处理批次中的每个文档
+                updated_documents_in_batch = [self._process_document_update(doc, metadata_dict, action) for doc in documents]
+                
+                # 6. 【Upsert逻辑】: 
+                if updated_documents_in_batch:
+                    self.logger.debug(f"Upserting batch of {len(updated_documents_in_batch)} documents using direct client upsert.")
+                    # 从更新后的 Document 对象中提取元数据字典列表
+                    upsert_data = [doc.metadata for doc in updated_documents_in_batch]
+                    self.client.upsert(
+                        collection_name=self.collection_name,
+                        data=upsert_data
+                    )
+                    total_updated_documents.extend(updated_documents_in_batch)
+
+        finally:
+            # 7. 【官方用法】: 确保迭代器被关闭
+            self.logger.info("Closing iterator.")
+            iterator.close()
+            
+        self.logger.info(f"Iterator processing complete. Total batches processed: {batch_i}.")
+        return total_updated_documents
+
+
+    def update_documents_metadata_folder_path(self, old_expr: str, metadata: str, action:Action=Action.UPSERT) -> List[Document]:
+        """
+        专门用于更新 version_metadata.folder_path 字段的方法。
+
+        它执行一个“目录移动”逻辑：
+        1. 使用 old_expr 找出所有路径匹配的文档。
+        2. 从 metadata 中获取新的基础路径。
+        3. 将文档中 folder_path 的 old_expr 前缀替换为新的基础路径，并保留后续的子路径。
+        """
+        # 1. 根据 old_expr 构造一个 "starts with" 查询
+        # Milvus JSON 'like' 操作符需要转义内部的双引号
+        # 但由于我们这里是变量，直接用 f-string 插入是安全的
+        expr = f"version_metadata[\"folder_path\"] like \"{old_expr}%\""
+
+        try:
+            metadata_dict = json.loads(metadata)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON string for metadata")
+
+        fields = self.get_fields() or []
+        # 确保关键字段都被查询出来
+        required_fields = ['pk', getattr(self, "_text_field", "text")]
+        vector_fields = self._vector_field if isinstance(self._vector_field, list) else [self._vector_field]
+        required_fields.extend(vector_fields)
+        
+        for f in required_fields:
+            if f not in fields:
+                fields.append(f)
+
+        total_updated_documents = []
+        batch_size = 1000
+
+        self.logger.info(f"Starting folder path update using 'collection.query_iterator' with expr: {expr}")
+
+        # self.col.load() # 确保集合已加载
+
+        iterator = self.col.query_iterator(
+            batch_size=batch_size,
+            expr=expr,
+            output_fields=fields
+        )
+
+        batch_i = 0
+        try:
+            while True:
+                batch_results = iterator.next()
+                if not batch_results:
+                    break
+
+                batch_i += 1
+                self.logger.info(f"Processing batch {batch_i} of {len(batch_results)} documents for folder path update.")
+
+                documents = [
+                    Document(page_content=result.get(getattr(self, "_text_field", "text"), ""), metadata=result)
+                    for result in batch_results
+                ]
+
+                updated_documents_in_batch = []
+                for doc in documents:
+                    # 沿用标准的 UPSERT 逻辑，但在处理 folder_path 时应用特殊规则
+                    if action == Action.UPSERT:
+                        existing_metadata = doc.metadata
+                        for key, value in metadata_dict.items():
+                            # ... (此处省略了标准的深层合并逻辑，与您已有的 update_documents_metadata 方法一致)
+                            # 仅展示与 folder_path 相关的特殊处理部分
+                            if isinstance(value, dict):
+                                # ... (处理从数据库读出的可能是字符串的JSON)
+                                if key in existing_metadata and isinstance(existing_metadata[key], str):
+                                    try:
+                                        existing_metadata[key] = json.loads(existing_metadata[key])
+                                    except json.JSONDecodeError:
+                                        existing_metadata[key] = {}
+                                
+                                if key not in existing_metadata or not isinstance(existing_metadata[key], dict):
+                                    existing_metadata[key] = value
+                                else:
+                                    # 递归更新，在这里注入我们的特殊逻辑
+                                    for sub_key, sub_value in value.items():
+                                        # 【核心特殊逻辑】
+                                        if key == 'version_metadata' and sub_key == 'folder_path':
+                                            new_folder_path_base = sub_value
+                                            current_folder_path = existing_metadata.get(key, {}).get(sub_key)
+
+                                            if current_folder_path and current_folder_path.startswith(old_expr):
+                                                # 移除旧前缀，保留子路径
+                                                sub_path = current_folder_path[len(old_expr):]
+                                                # 拼接新路径（确保斜杠正确）
+                                                new_full_path = f"{new_folder_path_base.rstrip('/')}/{sub_path.lstrip('/')}"
+                                                existing_metadata[key][sub_key] = new_full_path
+                                                self.logger.debug(f"Rewrote folder path from '{current_folder_path}' to '{new_full_path}'")
+                                            else:
+                                                # 如果不匹配，则按普通逻辑直接覆盖
+                                                existing_metadata[key][sub_key] = new_folder_path_base
+                                        
+                                        # 其他所有字段按原逻辑递归更新
+                                        elif isinstance(sub_value, dict):
+                                            if sub_key not in existing_metadata[key]:
+                                                existing_metadata[key][sub_key] = sub_value
+                                            else:
+                                                existing_metadata[key][sub_key].update(sub_value)
+                                        else:
+                                            existing_metadata[key][sub_key] = sub_value
+                            else:
+                                existing_metadata[key] = value
+                        doc.metadata = existing_metadata
+
+                    # (此处可以添加对 DELETE, UPDATE, INSERT 的处理，如果需要的话)
+                    
+                    updated_documents_in_batch.append(doc)
+
+                if updated_documents_in_batch:
+                    self.logger.debug(f"Upserting batch of {len(updated_documents_in_batch)} documents with updated folder paths.")
+                    upsert_data = [d.metadata for d in updated_documents_in_batch]
+                    self.client.upsert(
+                        collection_name=self.collection_name,
+                        data=upsert_data
+                    )
+                    total_updated_documents.extend(updated_documents_in_batch)
+        finally:
+            self.logger.info("Closing folder path update iterator.")
+            iterator.close()
+
+        self.logger.info(f"Folder path update complete. Total batches processed: {batch_i}.")
+        return total_updated_documents

{crewplus-0.2.39 → crewplus-0.2.41}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "crewplus"
-version = "0.2.39"
+version = "0.2.41"
 description = "Base services for CrewPlus AI applications"
 authors = [
     { name = "Tim Liu", email = "tim@opsmateai.com" },

crewplus-0.2.39/crewplus/vectorstores/milvus/schema_milvus.py

@@ -1,310 +0,0 @@
-from typing import List, Optional
-import logging
-import json
-
-from pymilvus import DataType
-from langchain_milvus import Milvus
-from langchain_core.documents import Document
-from ...utils.schema_document_updater import SchemaDocumentUpdater
-from ...utils.schema_action import Action
-from .milvus_schema_manager import MilvusSchemaManager
-
-DEFAULT_SCHEMA = """
-{
-    "node_types": {
-        "Document": {
-            "properties": {
-                "pk": {
-                    "type": "INT64",
-                    "is_primary": true,
-                    "auto_id": true
-                },
-                "vector": {
-                    "type": "FLOAT_VECTOR",
-                    "dim": 1536
-                },
-                "text": {
-                    "type": "VARCHAR",
-                    "max_length": 65535,
-                    "description": "The core text of the memory. This could be a user query, a documented fact, a procedural step, or a log of an event."
-                }
-	        }
-	    }
-    }
-}
-"""
-
-class SchemaMilvus(Milvus):
-    """
-    SchemaMilvus is a subclass of the Milvus class from langchain_milvus. This class is responsible for updating metadata of documents in a Milvus vector store.
-
-    Attributes:
-        embedding_function: Embedding function used by the Milvus vector store.
-        collection_name: Name of the collection in the Milvus vector store.
-        connection_args: Connection arguments for the Milvus vector store.
-        index_params: Index parameters for the Milvus vector store.
-        auto_id: Flag to specify if auto ID generation is enabled.
-        primary_field: The primary field of the collection.
-        vector_field: The vector field of the collection.
-        consistency_level: The consistency level for the Milvus vector store.
-        collection_schema: Schema JSON string associated with the Milvus existing collection name.
-    """
-    def __init__(
-        self, 
-        embedding_function, 
-        collection_name, 
-        connection_args, 
-        index_params=None, 
-        auto_id=True, 
-        primary_field="pk", 
-        text_field: str = "text",
-        vector_field=["vector"], 
-        consistency_level="Session",
-        logger: Optional[logging.Logger] = None
-    ):
-        """
-        Initializes the SchemaMilvus class with the provided parameters.
-
-        Args:
-            embedding_function: Embedding function used by the Milvus vector store.
-            collection_name: Name of the collection in the Milvus vector store.
-            connection_args: Connection arguments for the Milvus vector store.
-            index_params: Index parameters for the Milvus vector store.
-            auto_id: Flag to specify if auto ID generation is enabled.
-            primary_field: The primary field of the collection.
-            text_field: The text field of the collection.
-            vector_field: The vector field of the collection.
-            consistency_level: The consistency level for the Milvus vector store.
-            logger: Optional logger instance. If not provided, a default logger is created.
-        """
-        super().__init__(
-            embedding_function=embedding_function,
-            collection_name=collection_name,
-            connection_args=connection_args,
-            index_params=index_params,
-            auto_id=auto_id,
-            primary_field=primary_field,
-            text_field=text_field,
-            vector_field=vector_field,
-            consistency_level=consistency_level
-        )
-        self.logger = logger or logging.getLogger(__name__)
-        self.collection_schema = None
-        self.schema_manager = MilvusSchemaManager(client=self.client)
-
-    def set_schema(self, schema: str):
-        """
-        Sets the collection schema.
-
-        Args:
-            schema: The schema JSON string.
-        """
-        self.collection_schema = schema
-    
-    def get_fields(self, collection_name: Optional[str] = None) -> Optional[List[str]]:
-        """
-        Retrieves and returns the fields from the collection schema.
-
-        Args:
-            collection_name: The name of the collection to describe. If None, use self.collection_name.
-
-        Returns:
-            List[str] | None: The list of field names from the collection schema (excluding vector and text fields), or None if collection_name is not provided or an error occurs.
-        """
-        if collection_name is None:
-            collection_name = self.collection_name
-        if collection_name is None:
-            return None
-
-        try:
-            schema = self.client.describe_collection(collection_name)
-            fields = [field["name"] for field in schema["fields"] if field["type"] != DataType.FLOAT_VECTOR ]
-            return fields
-        except Exception as e:
-            self.logger.warning(f"Failed to retrieve schema fields: {e}")
-            return None
-    
-    def create_collection(self) -> bool:
-        """
-        Validates the schema and creates the collection using the MilvusSchemaManager.
-
-        Returns:
-            bool: True if the collection is successfully created, False otherwise.
-        """
-        if self.collection_schema is None:
-            self.logger.error("Collection schema is not set. Please set a schema using set_schema().")
-            return False
-            
-        self.schema_manager.bind_client(self.client)
-        if not self.schema_manager.validate_schema(self.collection_schema):
-            self.logger.error("Failed to validate schema")
-            return False
-        try:
-            self.schema_manager.create_collection(self.collection_name, self.collection_schema)
-            self.logger.info(f"Collection {self.collection_name} created successfully")
-            
-            return True
-        except Exception as e:
-            self.logger.error(f"Failed to create collection: {e}")
-            return False
-
-    def drop_collection(self, collection_name: Optional[str] = None) -> bool:
-        """
-        Drops the collection using the Milvus client.
-
-        Returns:
-            bool: True if the collection is successfully dropped, False otherwise.
-        """
-        if collection_name is None:
-            collection_name = self.collection_name
-            
-        try:
-            self.client.drop_collection(collection_name)
-            self.logger.info(f"Collection {collection_name} dropped successfully")
-            return True
-        except Exception as e:
-            self.logger.error(f"Failed to drop collection {self.collection_name}: {e}")
-            return False
-
-    def _handle_upsert(self, doc: Document, metadata_dict: dict) -> Document:
-        """
-        Handles the UPSERT action for a single document by merging metadata.
-        """
-        existing_metadata = doc.metadata
-        for key, value in metadata_dict.items():
-            # Skip primary key and text fields to prevent modification.
-            if key in [self.primary_field, self.text_field]:
-                continue
-
-            if isinstance(value, dict):
-                # If the new value is a dictionary, handle nested updates.
-                if key not in existing_metadata or not isinstance(existing_metadata.get(key), dict):
-                    # If the key doesn't exist or its value is not a dict, replace it.
-                    existing_metadata[key] = value
-                else:
-                    # If both are dictionaries, recursively update the nested fields.
-                    for sub_key, sub_value in value.items():
-                        if isinstance(sub_value, dict) and sub_key in existing_metadata[key] and isinstance(existing_metadata[key].get(sub_key), dict):
-                            existing_metadata[key][sub_key].update(sub_value)
-                        else:
-                            existing_metadata[key][sub_key] = sub_value
-            else:
-                # For non-dictionary values, simply update or add the field.
-                existing_metadata[key] = value
-
-        doc.metadata = existing_metadata
-        return doc
-
-    def _prepare_documents_for_update(self, expr: str, metadata: str, action: Action = Action.UPSERT) -> tuple[List[Document], List]:
-        try:
-            metadata_dict = json.loads(metadata)
-        except json.JSONDecodeError:
-            raise ValueError("Invalid JSON string for metadata")
-
-        # Retrieve documents that match the filter expression.
-        fields = self.get_fields()
-        documents = self.search_by_metadata(expr, fields=fields, limit=5000)
-
-        updated_documents = []
-        for doc in documents:
-            # Preserve the original primary key and text values.
-            pk_value = doc.metadata.get(self.primary_field) # default to pk
-            text_value = doc.metadata.get(self.text_field)
-
-            # Apply the specified action to update the document's metadata.
-            if action == Action.UPSERT:
-                doc = self._handle_upsert(doc, metadata_dict)
-            elif action == Action.DELETE:
-                keys_to_delete = metadata_dict.keys()
-                doc = SchemaDocumentUpdater.delete_document_metadata(doc, list(keys_to_delete))
-            elif action == Action.UPDATE:
-                existing_metadata = doc.metadata
-                update_dict = {}
-                for key, value in metadata_dict.items():
-                    if key in existing_metadata:
-                        if isinstance(value, dict) and isinstance(existing_metadata.get(key), dict):
-                            merged = existing_metadata[key].copy()
-                            for sub_key, sub_value in value.items():
-                                if sub_key in merged:
-                                    merged[sub_key] = sub_value
-                            update_dict[key] = merged
-                        else:
-                            update_dict[key] = value
-                doc = SchemaDocumentUpdater.update_document_metadata(doc, update_dict)
-            elif action == Action.INSERT:
-                existing_metadata = doc.metadata
-                for key, value in metadata_dict.items():
-                    if key in ['pk', 'text']:
-                        continue
-
-                    if isinstance(value, dict) and key in existing_metadata and isinstance(existing_metadata.get(key), dict):
-                        existing_metadata[key] = {}
-                        existing_metadata[key] = value
-                    else:
-                        existing_metadata[key] = value
-                doc.metadata = existing_metadata
-
-            # Restore the primary key and text values to ensure they are not lost.
-            if pk_value is not None:
-                doc.metadata[self.primary_field] = pk_value
-            if text_value is not None:
-                doc.metadata[self.text_field] = text_value
-
-            updated_documents.append(doc)
-
-        if not updated_documents:
-            return [], []
-
-        # Extract the primary keys for the upsert operation.
-        updated_ids = [doc.metadata[self.primary_field] for doc in updated_documents]
-
-        # Remove primary key and text from metadata before upserting,
-        # as they are handled separately by the vector store.
-        for doc in updated_documents:
-            doc.metadata.pop(self.primary_field, None)
-            doc.metadata.pop(self.text_field, None)
-
-        return updated_documents, updated_ids
-
-    def update_documents_metadata(self, expr: str, metadata: str, action: Action = Action.UPSERT) -> List[Document]:
-        """
-        Updates the metadata of documents in the Milvus vector store based on the provided expression.
-
-        Args:
-            expr: Expression to filter the target documents.
-            metadata: New metadata to update the documents with.
-
-        Returns:
-            List of updated documents.
-        """
-        documents_to_upsert, updated_ids = self._prepare_documents_for_update(expr, metadata, action)
-
-        if not documents_to_upsert:
-            return []
-
-        # Perform the upsert operation to update the documents in the collection.
-        self.upsert(ids=updated_ids, documents=documents_to_upsert)
-
-        return documents_to_upsert
-
-    async def aupdate_documents_metadata(self, expr: str, metadata: str, action: Action = Action.UPSERT) -> List[Document]:
-        """
-        Asynchronously updates the metadata of documents in the Milvus vector store.
-
-        Args:
-            expr: Expression to filter the target documents.
-            metadata: New metadata to update the documents with.
-            action: The action to perform on the document metadata.
-
-        Returns:
-            List of updated documents.
-        """
-        documents_to_upsert, updated_ids = self._prepare_documents_for_update(expr, metadata, action)
-
-        if not documents_to_upsert:
-            return []
-
-        # Perform the asynchronous upsert operation.
-        await self.aupsert(ids=updated_ids, documents=documents_to_upsert)
-
-        return documents_to_upsert