Merge pull request #10 from FlowLLM-AI/vector_store_update

refactor(core): rename async embedding methods and update vector store

Author: jinliyl

.github/workflows/pre-commit.yml

@@ -0,0 +1,38 @@
+name: Pre-commit
+
+on: [ push, pull_request ]
+
+jobs:
+  run:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: True
+      matrix:
+        os: [ ubuntu-latest ]
+    env:
+      OS: ${{ matrix.os }}
+      PYTHON: '3.10'
+    steps:
+      - uses: actions/checkout@master
+      - name: Setup Python
+        uses: actions/setup-python@master
+        with:
+          python-version: '3.10'
+      - name: Update setuptools
+        run: |
+          pip install -U setuptools wheel
+      - name: Install
+        run: |
+          pip install -q -e .[dev]
+      - name: Install pre-commit
+        run: |
+          pre-commit install
+      - name: Pre-commit starts
+        run: |
+          pre-commit run --all-files > pre-commit.log 2>&1 || true
+          cat pre-commit.log
+          if grep -q Failed pre-commit.log; then
+            echo -e "\e[41m  [**FAIL**] Please install pre-commit and format your code first. \e[0m"
+            exit 1
+          fi
+          echo -e "\e[46m  ********************************Passed******************************** \e[0m"

docs/zh/guide/async_op_advance_guide.md

@@ -30,7 +30,7 @@ class RAGSearchOp(BaseAsyncOp):
         assert query, "query 不能为空"
 
         # 2. 使用 Embedding 生成查询向量（一行代码示意 emb）
-        query_vector = await self.embedding_model.get_embeddings_async(query)
+        query_vector = await self.embedding_model.async_get_embeddings(query)
 
         # 3. 使用 VectorStore 进行语义搜索（一行代码示意 vectorstore）
         nodes = await self.vector_store.async_search(query=query, workspace_id=workspace_id, top_k=top_k)

flowllm/__init__.py

@@ -11,4 +11,4 @@
 from . import gallery  # noqa: E402, F401  # pylint: disable=wrong-import-position,unused-import
 from . import extensions  # noqa: E402, F401  # pylint: disable=wrong-import-position,unused-import
 
-__version__ = "0.2.0.8"
+__version__ = "0.2.0.9"

flowllm/core/embedding_model/base_embedding_model.py

@@ -62,7 +62,7 @@ def _get_embeddings(self, input_text: str | List[str]):
         """
         raise NotImplementedError
 
-    async def _get_embeddings_async(self, input_text: str | List[str]):
+    async def _async_get_embeddings(self, input_text: str | List[str]):
         """
         Abstract async method to get embeddings from the model.
 
@@ -104,7 +104,7 @@ def get_embeddings(self, input_text: str | List[str]):
         # Return None if all retries failed and raise_exception is False
         return None
 
-    async def get_embeddings_async(self, input_text: str | List[str]):
+    async def async_get_embeddings(self, input_text: str | List[str]):
         """
         Get embeddings asynchronously with retry logic and error handling.
 
@@ -120,7 +120,7 @@ async def get_embeddings_async(self, input_text: str | List[str]):
         # Retry loop with exponential backoff potential
         for i in range(self.max_retries):
             try:
-                return await self._get_embeddings_async(input_text)
+                return await self._async_get_embeddings(input_text)
 
             except Exception as e:
                 logger.exception(f"embedding model name={self.model_name} encounter error with e={e.args}")
@@ -173,7 +173,7 @@ def get_node_embeddings(self, nodes: VectorNode | List[VectorNode]):
         else:
             raise TypeError(f"unsupported type={type(nodes)}")
 
-    async def get_node_embeddings_async(self, nodes: VectorNode | List[VectorNode]):
+    async def async_get_node_embeddings(self, nodes: VectorNode | List[VectorNode]):
         """
         Generate embeddings asynchronously for VectorNode objects and update their vector fields.
 
@@ -191,7 +191,7 @@ async def get_node_embeddings_async(self, nodes: VectorNode | List[VectorNode]):
         """
         # Handle single VectorNode
         if isinstance(nodes, VectorNode):
-            nodes.vector = await self.get_embeddings_async(nodes.content)
+            nodes.vector = await self.async_get_embeddings(nodes.content)
             return nodes
 
         # Handle list of VectorNodes with batch processing
@@ -201,7 +201,7 @@ async def get_node_embeddings_async(self, nodes: VectorNode | List[VectorNode]):
             for i in range(0, len(nodes), self.max_batch_size):
                 batch_nodes = nodes[i : i + self.max_batch_size]
                 batch_content = [node.content for node in batch_nodes]
-                batch_tasks.append(self.get_embeddings_async(batch_content))
+                batch_tasks.append(self.async_get_embeddings(batch_content))
 
             # Execute all batch tasks concurrently
             batch_results = await asyncio.gather(*batch_tasks)
@@ -220,3 +220,9 @@ async def get_node_embeddings_async(self, nodes: VectorNode | List[VectorNode]):
 
         else:
             raise TypeError(f"unsupported type={type(nodes)}")
+
+    def close(self):
+        """Close the client connection or clean up resources."""
+
+    async def async_close(self):
+        """Asynchronously close the client connection or clean up resources."""

flowllm/core/embedding_model/openai_compatible_embedding_model.py

@@ -105,7 +105,7 @@ def _get_embeddings(self, input_text: str | List[str]):
         else:
             raise RuntimeError(f"unsupported type={type(input_text)}")
 
-    async def _get_embeddings_async(self, input_text: str | List[str]):
+    async def _async_get_embeddings(self, input_text: str | List[str]):
         """
         Get embeddings asynchronously from the OpenAI-compatible API.
 
@@ -139,3 +139,11 @@ async def _get_embeddings_async(self, input_text: str | List[str]):
 
         else:
             raise RuntimeError(f"unsupported type={type(input_text)}")
+
+    def close(self):
+        """Close the OpenAI clients and clean up resources."""
+        self._client.close()
+
+    async def async_close(self):
+        """Asynchronously close the OpenAI clients and clean up resources."""
+        await self._async_client.close()

flowllm/core/llm/base_llm.py

@@ -260,3 +260,9 @@ async def achat(
                 await asyncio.sleep(1 + i)
 
         return default_value
+
+    def close(self):
+        """Close the client connection or clean up resources."""
+
+    async def async_close(self):
+        """Asynchronously close the client connection or clean up resources."""

flowllm/core/llm/openai_compatible_llm.py

@@ -429,3 +429,9 @@ async def _achat(
             content=answer_content,
             tool_calls=tool_calls,
         )
+
+    def close(self):
+        self._client.close()
+
+    async def async_close(self):
+        await self._aclient.close()