flowllm 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

llmflow/vector_store/file_vector_store.py → flowllm/storage/vector_store/local_vector_store.py

@@ -1,18 +1,20 @@
+import fcntl
+import json
 import math
 from pathlib import Path
 from typing import List, Iterable
 
 from loguru import logger
 from pydantic import Field, model_validator
+from tqdm import tqdm
 
-from llmflow.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
-from llmflow.schema.vector_node import VectorNode
-from llmflow.vector_store import VECTOR_STORE_REGISTRY
-from llmflow.vector_store.base_vector_store import BaseVectorStore
+from flowllm.context.service_context import C
+from flowllm.schema.vector_node import VectorNode
+from flowllm.storage.vector_store.base_vector_store import BaseVectorStore
 
 
-@VECTOR_STORE_REGISTRY.register("local_file")
-class FileVectorStore(BaseVectorStore):
+@C.register_vector_store("local")
+class LocalVectorStore(BaseVectorStore):
     store_dir: str = Field(default="./file_vector_store")
 
     @model_validator(mode="after")
@@ -21,6 +23,55 @@ class FileVectorStore(BaseVectorStore):
         store_path.mkdir(parents=True, exist_ok=True)
         return self
 
+    @staticmethod
+    def _load_from_path(workspace_id: str, path: str | Path, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
+        workspace_path = Path(path) / f"{workspace_id}.jsonl"
+        if not workspace_path.exists():
+            logger.warning(f"workspace_path={workspace_path} is not exists!")
+            return
+
+        with workspace_path.open() as f:
+            fcntl.flock(f, fcntl.LOCK_SH)
+            try:
+                for line in tqdm(f, desc="load from path"):
+                    if line.strip():
+                        node_dict = json.loads(line.strip())
+                        if callback_fn:
+                            node = callback_fn(node_dict)
+                        else:
+                            node = VectorNode(**node_dict, **kwargs)
+                        node.workspace_id = workspace_id
+                        yield node
+
+            finally:
+                fcntl.flock(f, fcntl.LOCK_UN)
+
+    @staticmethod
+    def _dump_to_path(nodes: Iterable[VectorNode], workspace_id: str, path: str | Path = "", callback_fn=None,
+                      ensure_ascii: bool = False, **kwargs):
+        dump_path: Path = Path(path)
+        dump_path.mkdir(parents=True, exist_ok=True)
+        dump_file = dump_path / f"{workspace_id}.jsonl"
+
+        count = 0
+        with dump_file.open("w") as f:
+            fcntl.flock(f, fcntl.LOCK_EX)
+            try:
+                for node in tqdm(nodes, desc="dump to path"):
+                    node.workspace_id = workspace_id
+                    if callback_fn:
+                        node_dict = callback_fn(node)
+                    else:
+                        node_dict = node.model_dump()
+                    assert isinstance(node_dict, dict)
+                    f.write(json.dumps(node_dict, ensure_ascii=ensure_ascii, **kwargs))
+                    f.write("\n")
+                    count += 1
+
+                return {"size": count}
+            finally:
+                fcntl.flock(f, fcntl.LOCK_UN)
+
     @property
     def store_path(self) -> Path:
         return Path(self.store_dir)
@@ -41,6 +92,54 @@ class FileVectorStore(BaseVectorStore):
         for i, node in enumerate(self._load_from_path(path=self.store_path, workspace_id=workspace_id, **kwargs)):
             yield node
 
+    def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
+        if not self.exist_workspace(workspace_id=workspace_id, **kwargs):
+            logger.warning(f"workspace_id={workspace_id} is not exist!")
+            return {}
+
+        return self._dump_to_path(nodes=self._iter_workspace_nodes(workspace_id=workspace_id, **kwargs),
+                                  workspace_id=workspace_id,
+                                  path=path,
+                                  callback_fn=callback_fn,
+                                  **kwargs)
+
+    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, callback_fn=None,
+                       **kwargs):
+        if self.exist_workspace(workspace_id, **kwargs):
+            self.delete_workspace(workspace_id=workspace_id, **kwargs)
+            logger.info(f"delete workspace_id={workspace_id}")
+
+        self.create_workspace(workspace_id=workspace_id, **kwargs)
+
+        all_nodes: List[VectorNode] = []
+        if nodes:
+            all_nodes.extend(nodes)
+        for node in self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs):
+            all_nodes.append(node)
+        self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs)
+        return {"size": len(all_nodes)}
+
+    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
+        if not self.exist_workspace(workspace_id=src_workspace_id, **kwargs):
+            logger.warning(f"src_workspace_id={src_workspace_id} is not exist!")
+            return {}
+
+        if not self.exist_workspace(dest_workspace_id, **kwargs):
+            self.create_workspace(workspace_id=dest_workspace_id, **kwargs)
+
+        nodes = []
+        node_size = 0
+        for node in self._iter_workspace_nodes(workspace_id=src_workspace_id, **kwargs):
+            nodes.append(node)
+            node_size += 1
+            if len(nodes) >= self.batch_size:
+                self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs)
+                nodes.clear()
+
+        if nodes:
+            self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs)
+        return {"size": node_size}
+
     @staticmethod
     def calculate_similarity(query_vector: List[float], node_vector: List[float]):
         assert query_vector, f"query_vector is empty!"
@@ -106,12 +205,14 @@ class FileVectorStore(BaseVectorStore):
 
 
 def main():
-    from dotenv import load_dotenv
-    load_dotenv()
+    from flowllm.utils.common_utils import load_env
+    from flowllm.embedding_model import OpenAICompatibleEmbeddingModel
+
+    load_env()
 
     embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
     workspace_id = "rag_nodes_index"
-    client = FileVectorStore(embedding_model=embedding_model)
+    client = LocalVectorStore(embedding_model=embedding_model)
     client.delete_workspace(workspace_id)
     client.create_workspace(workspace_id)
 
@@ -160,4 +261,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-    # launch with: python -m llmflow.storage.file_vector_store

flowllm/utils/common_utils.py

@@ -0,0 +1,64 @@
+import os
+import re
+from pathlib import Path
+
+from loguru import logger
+
+
+def camel_to_snake(content: str) -> str:
+    """
+    BaseWorker -> base_worker
+    """
+    snake_str = re.sub(r'(?<!^)(?=[A-Z])', '_', content).lower()
+    return snake_str
+
+
+def snake_to_camel(content: str) -> str:
+    """
+    base_worker -> BaseWorker
+    """
+    camel_str = "".join(x.capitalize() for x in content.split("_"))
+    return camel_str
+
+
+def _load_env(path: Path):
+    with path.open() as f:
+        for line in f:
+            line = line.strip()
+            if line.startswith("#"):
+                continue
+
+            line_split = line.strip().split("=", 1)
+            if len(line_split) >= 2:
+                key = line_split[0].strip()
+                value = line_split[1].strip()
+                os.environ[key] = value
+
+
+def load_env(path: str | Path = None):
+    if path is not None:
+        path = Path(path)
+        if path.exists():
+            _load_env(path)
+    else:
+        path1 = Path(".env")
+        path2 = Path("../.env")
+        path3 = Path("../../.env")
+        path4 = Path("../../../.env")
+        path5 = Path("../../../.env")
+
+        if path1.exists():
+            path = path1
+        elif path2.exists():
+            path = path2
+        elif path3.exists():
+            path = path3
+        elif path4.exists():
+            path = path4
+        elif path5.exists():
+            path = path5
+        else:
+            raise FileNotFoundError(".env not found")
+
+        logger.info(f"using path={path}")
+        _load_env(path)

flowllm/utils/dataframe_cache.py

@@ -0,0 +1,331 @@
+"""
+DataFrame cache utility that supports local CSV file storage and reading with data expiration functionality
+"""
+
+import json
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Optional, Dict, Any
+
+import pandas as pd
+
+
+class DataFrameCache:
+    """
+    DataFrame cache utility class
+    
+    Features:
+    - Support for pandas DataFrame local CSV storage and reading
+    - Support for data expiration time settings
+    - Automatic cleanup of expired data
+    - Recording and managing update timestamps
+    """
+
+    def __init__(self, cache_dir: str = "cache_df"):
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.metadata_file = self.cache_dir / "metadata.json"
+        self.metadata = {}
+        self._load_metadata()
+
+    def _load_metadata(self):
+        """Load metadata"""
+        if self.metadata_file.exists():
+            with open(self.metadata_file) as f:
+                self.metadata = json.load(f)
+
+    def _save_metadata(self):
+        """Save metadata"""
+        with open(self.metadata_file, "w") as f:
+            json.dump(self.metadata, f, ensure_ascii=False, indent=2)
+
+    def _get_file_path(self, key: str) -> Path:
+        """Get data file path"""
+        return self.cache_dir / f"{key}.csv"
+
+    def _is_expired(self, key: str) -> bool:
+        """Check if data is expired"""
+        if key not in self.metadata:
+            return True
+
+        expire_time_str = self.metadata[key].get('expire_time')
+        if not expire_time_str:
+            return False  # No expiration time set, never expires
+
+        expire_time = datetime.fromisoformat(expire_time_str)
+        return datetime.now() > expire_time
+
+    def save(self, key: str, df: pd.DataFrame, expire_hours: Optional[float] = None,
+             **csv_kwargs) -> bool:
+        """
+        Save DataFrame to cache
+        
+        Args:
+            key: Cache key name
+            df: DataFrame to save
+            expire_hours: Expiration time in hours, None means never expires
+            **csv_kwargs: Additional parameters passed to pandas to_csv
+            
+        Returns:
+            bool: Whether save was successful
+        """
+        try:
+            file_path = self._get_file_path(key)
+
+            # Set default CSV parameters
+            csv_params = {
+                "index": False,
+                "encoding": "utf-8"
+            }
+            csv_params.update(csv_kwargs)
+
+            # Save CSV file
+            df.to_csv(file_path, **csv_params)
+
+            # Update metadata
+            current_time = datetime.now()
+            self.metadata[key] = {
+                'created_time': current_time.isoformat(),
+                'updated_time': current_time.isoformat(),
+                'expire_time': (current_time + timedelta(hours=expire_hours)).isoformat() if expire_hours else None,
+                'file_size': file_path.stat().st_size,
+                'row_count': len(df),
+                'column_count': len(df.columns)
+            }
+
+            self._save_metadata()
+            return True
+
+        except Exception as e:
+            print(f"Failed to save DataFrame: {e}")
+            return False
+
+    def load(self, key: str, auto_clean_expired: bool = True, **csv_kwargs) -> Optional[pd.DataFrame]:
+        """
+        Load DataFrame from cache
+        
+        Args:
+            key: Cache key name
+            auto_clean_expired: Whether to automatically clean expired data
+            **csv_kwargs: Additional parameters passed to pandas read_csv
+            
+        Returns:
+            Optional[pd.DataFrame]: Loaded DataFrame, returns None if not exists or expired
+        """
+        try:
+            # Check if expired
+            if self._is_expired(key):
+                if auto_clean_expired:
+                    self.delete(key)
+                    print(f"Cache '{key}' has expired and was automatically cleaned")
+                return None
+
+            file_path = self._get_file_path(key)
+            if not file_path.exists():
+                return None
+
+            # Set default CSV parameters
+            csv_params = {
+                'encoding': 'utf-8'
+            }
+            csv_params.update(csv_kwargs)
+
+            # Read CSV file
+            df = pd.read_csv(file_path, **csv_params)
+
+            # Update last access time
+            if key in self.metadata:
+                self.metadata[key]['last_accessed'] = datetime.now().isoformat()
+                self._save_metadata()
+
+            return df
+
+        except Exception as e:
+            print(f"Failed to load DataFrame: {e}")
+            return None
+
+    def exists(self, key: str, check_expired: bool = True) -> bool:
+        """
+        Check if cache exists
+        
+        Args:
+            key: Cache key name
+            check_expired: Whether to check expiration status
+            
+        Returns:
+            bool: Whether cache exists and is not expired
+        """
+        if check_expired and self._is_expired(key):
+            return False
+
+        file_path = self._get_file_path(key)
+        return file_path.exists() and key in self.metadata
+
+    def delete(self, key: str) -> bool:
+        """
+        Delete cache
+        
+        Args:
+            key: Cache key name
+            
+        Returns:
+            bool: Whether deletion was successful
+        """
+        try:
+            file_path = self._get_file_path(key)
+
+            # Delete CSV file
+            if file_path.exists():
+                file_path.unlink()
+
+            # Delete metadata
+            if key in self.metadata:
+                del self.metadata[key]
+                self._save_metadata()
+
+            return True
+
+        except Exception as e:
+            print(f"Failed to delete cache: {e}")
+            return False
+
+    def clean_expired(self) -> int:
+        """
+        Clean all expired caches
+        
+        Returns:
+            int: Number of cleaned caches
+        """
+        expired_keys = []
+
+        for key in list(self.metadata.keys()):
+            if self._is_expired(key):
+                expired_keys.append(key)
+
+        cleaned_count = 0
+        for key in expired_keys:
+            if self.delete(key):
+                cleaned_count += 1
+
+        return cleaned_count
+
+    def get_info(self, key: str) -> Optional[Dict[str, Any]]:
+        """
+        Get cache information
+        
+        Args:
+            key: Cache key name
+            
+        Returns:
+            Optional[Dict]: Cache information including creation time, update time, expiration time, etc.
+        """
+        if key not in self.metadata:
+            return None
+
+        info = self.metadata[key].copy()
+        info['key'] = key
+        info['is_expired'] = self._is_expired(key)
+        info['file_path'] = str(self._get_file_path(key))
+
+        return info
+
+    def list_all(self, include_expired: bool = False) -> Dict[str, Dict[str, Any]]:
+        """
+        List all caches
+        
+        Args:
+            include_expired: Whether to include expired caches
+            
+        Returns:
+            Dict: Information of all caches
+        """
+        result = {}
+
+        for key in self.metadata:
+            if not include_expired and self._is_expired(key):
+                continue
+
+            info = self.get_info(key)
+            if info:
+                result[key] = info
+
+        return result
+
+    def get_cache_stats(self) -> Dict[str, Any]:
+        """
+        Get cache statistics
+        
+        Returns:
+            Dict: Cache statistics information
+        """
+        total_count = len(self.metadata)
+        expired_count = sum(1 for key in self.metadata if self._is_expired(key))
+        active_count = total_count - expired_count
+
+        total_size = 0
+        for key in self.metadata:
+            file_path = self._get_file_path(key)
+            if file_path.exists():
+                total_size += file_path.stat().st_size
+
+        return {
+            'total_count': total_count,
+            'active_count': active_count,
+            'expired_count': expired_count,
+            'total_size_bytes': total_size,
+            'total_size_mb': round(total_size / (1024 * 1024), 2),
+            'cache_dir': str(self.cache_dir)
+        }
+
+    def clear_all(self) -> bool:
+        """
+        Clear all caches
+        
+        Returns:
+            bool: Whether clearing was successful
+        """
+        try:
+            # Delete all CSV files
+            for csv_file in self.cache_dir.glob("*.csv"):
+                csv_file.unlink()
+
+            # Clear metadata
+            self.metadata = {}
+            self._save_metadata()
+
+            return True
+
+        except Exception as e:
+            print(f"Failed to clear cache: {e}")
+            return False
+
+
+# Create default instance
+default_cache = DataFrameCache()
+
+
+# Convenience functions
+def save_dataframe(key: str, df: pd.DataFrame, expire_hours: Optional[float] = None,
+                   **csv_kwargs) -> bool:
+    """Convenience function: Save DataFrame"""
+    return default_cache.save(key, df, expire_hours, **csv_kwargs)
+
+
+def load_dataframe(key: str, **csv_kwargs) -> Optional[pd.DataFrame]:
+    """Convenience function: Load DataFrame"""
+    return default_cache.load(key, **csv_kwargs)
+
+
+def dataframe_exists(key: str) -> bool:
+    """Convenience function: Check if DataFrame exists"""
+    return default_cache.exists(key)
+
+
+def delete_dataframe(key: str) -> bool:
+    """Convenience function: Delete DataFrame cache"""
+    return default_cache.delete(key)
+
+
+def clean_expired_dataframes() -> int:
+    """Convenience function: Clean expired DataFrame caches"""
+    return default_cache.clean_expired()

flowllm/utils/fetch_url.py

@@ -0,0 +1,113 @@
+import requests
+from bs4 import BeautifulSoup
+import urllib3
+import random
+import time
+
+# Disable SSL warnings (optional, for handling insecure HTTPS)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+def get_random_headers():
+    """Generate random headers to avoid detection"""
+    user_agents = [
+        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
+        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.107 Safari/537.36',
+        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/93.0.4577.82 Safari/537.36',
+        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
+        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.1.1 Safari/605.1.15',
+        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
+        'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0',
+        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:89.0) Gecko/20100101 Firefox/89.0',
+        'Mozilla/5.0 (X11; Linux x86_64; rv:89.0) Gecko/20100101 Firefox/89.0',
+        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Edge/91.0.864.59'
+    ]
+    
+    accept_languages = [
+        'en-US,en;q=0.9',
+        'zh-CN,zh;q=0.9,en;q=0.8',
+        'en-GB,en;q=0.9',
+        'fr-FR,fr;q=0.9,en;q=0.8',
+        'de-DE,de;q=0.9,en;q=0.8'
+    ]
+    
+    accept_encodings = [
+        'gzip, deflate, br',
+        'gzip, deflate',
+        'br, gzip, deflate'
+    ]
+    
+    headers = {
+        'User-Agent': random.choice(user_agents),
+        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
+        'Accept-Language': random.choice(accept_languages),
+        'Accept-Encoding': random.choice(accept_encodings),
+        'Connection': 'keep-alive',
+        'Upgrade-Insecure-Requests': '1',
+    }
+    
+    # Randomly add some optional headers
+    if random.random() > 0.5:
+        headers['DNT'] = '1'
+    if random.random() > 0.7:
+        headers['Cache-Control'] = 'max-age=0'
+    if random.random() > 0.6:
+        headers['Sec-Fetch-Dest'] = 'document'
+        headers['Sec-Fetch-Mode'] = 'navigate'
+        headers['Sec-Fetch-Site'] = 'none'
+    
+    return headers
+
+def fetch_webpage_text(url, min_delay=1, max_delay=3):
+    """
+    Fetch and extract text content from a webpage with randomization
+    
+    Args:
+        url (str): The URL to fetch
+        min_delay (int): Minimum delay before request (seconds)
+        max_delay (int): Maximum delay before request (seconds)
+    
+    Returns:
+        str: Extracted text content or error message
+    """
+    # Add random delay to avoid being detected as bot
+    delay = random.uniform(min_delay, max_delay)
+    time.sleep(delay)
+    
+    headers = get_random_headers()
+    
+    # Random timeout between 8-15 seconds
+    timeout = random.randint(8, 15)
+    
+    try:
+        # Send request with random headers and timeout
+        response = requests.get(url, headers=headers, timeout=timeout, verify=False)
+        response.raise_for_status()  # Check if request was successful
+        response.encoding = response.apparent_encoding  # Auto-detect encoding
+
+        # Parse HTML using BeautifulSoup
+        soup = BeautifulSoup(response.text, 'html.parser')
+
+        # Remove script, style and navigation elements to avoid interference
+        for script_or_style in soup(['script', 'style', 'nav', 'footer', 'header']):
+            script_or_style.decompose()
+
+        # Extract text content
+        text = soup.get_text()
+
+        # Clean whitespace: remove extra blank lines and spaces
+        lines = (line.strip() for line in text.splitlines())
+        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
+        text = ' '.join(chunk for chunk in chunks if chunk)
+
+        return text
+
+    except requests.exceptions.RequestException as e:
+        return f"Request failed: {e}"
+    except Exception as e:
+        return f"Parsing failed: {e}"
+
+# Example usage
+if __name__ == "__main__":
+    url = "http://finance.eastmoney.com/a/202508133482756869.html"
+    text = fetch_webpage_text(url)
+    print(text)

{llmflow → flowllm}/utils/timer.py

@@ -1,9 +1,10 @@
 import time
+from typing import Optional
 
 from loguru import logger
 
 
-class Timer(object):
+class Timer:
     def __init__(self, name: str, use_ms: bool = False, stack_level: int = 2):
         self.name: str = name
         self.use_ms: bool = use_ms
@@ -15,7 +16,7 @@ class Timer(object):
 
     def __enter__(self, *args, **kwargs):
         self.time_start = time.time()
-        logger.info(f"---------- enter {self.name} ----------", stacklevel=self.stack_level)
+        logger.info(f"========== timer.{self.name} start ==========", stacklevel=self.stack_level)
         return self
 
     def __exit__(self, *args):
@@ -26,10 +27,10 @@ class Timer(object):
         else:
             time_str = f"{self.time_cost:.3f}s"
 
-        logger.info(f"---------- leave {self.name} [{time_str}] ----------", stacklevel=self.stack_level)
+        logger.info(f"========== timer.{self.name} end, time_cost={time_str} ==========", stacklevel=self.stack_level)
 
 
-def timer(name: str = None, use_ms: bool = False, stack_level: int = 2):
+def timer(name: Optional[str] = None, use_ms: bool = False, stack_level: int = 2):
     def decorator(func):
         def wrapper(*args, **kwargs):
             with Timer(name=name or func.__name__, use_ms=use_ms, stack_level=stack_level + 1):