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

flowllm/{utils/dataframe_cache.py → storage/cache/data_cache.py}

@@ -1,33 +1,65 @@
 """
-DataFrame cache utility that supports local CSV file storage and reading with data expiration functionality
+DataCache utility that supports multiple data types with local storage and data expiration functionality
 """
 
 import json
 from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, Union, Type
 
 import pandas as pd
 
+from flowllm.storage.cache.cache_data_handler import CacheDataHandler, DataFrameHandler, DictHandler
 
-class DataFrameCache:
+
+class DataCache:
     """
-    DataFrame cache utility class
-    
+    Generic data cache utility class
+
     Features:
-    - Support for pandas DataFrame local CSV storage and reading
+    - Support for multiple data types (DataFrame, dict, and extensible for others)
     - Support for data expiration time settings
     - Automatic cleanup of expired data
     - Recording and managing update timestamps
+    - Type-specific storage formats (CSV for DataFrame, JSON for dict)
     """
 
-    def __init__(self, cache_dir: str = "cache_df"):
+    def __init__(self, cache_dir: str = "cache"):
         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 = {}
+
+        # Register default handlers
+        self.handlers: Dict[Type, CacheDataHandler] = {
+            pd.DataFrame: DataFrameHandler(),
+            dict: DictHandler()
+        }
+
         self._load_metadata()
 
+    def register_handler(self, data_type: Type, handler: CacheDataHandler):
+        """
+        Register a custom data handler for a specific data type
+
+        Args:
+            data_type: The data type to handle
+            handler: The handler instance
+        """
+        self.handlers[data_type] = handler
+
+    def _get_handler(self, data_type: Type) -> CacheDataHandler:
+        """Get the appropriate handler for a data type"""
+        if data_type in self.handlers:
+            return self.handlers[data_type]
+
+        # Try to find a handler for parent classes
+        for registered_type, handler in self.handlers.items():
+            if issubclass(data_type, registered_type):
+                return handler
+
+        raise ValueError(f"No handler registered for data type: {data_type}")
+
     def _load_metadata(self):
         """Load metadata"""
         if self.metadata_file.exists():
@@ -39,9 +71,35 @@ class DataFrameCache:
         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 _get_file_path(self, key: str, data_type: Type = None) -> Path:
+        """Get data file path with appropriate extension"""
+        if data_type is None:
+            # Try to get extension from metadata
+            if key in self.metadata and 'data_type' in self.metadata[key]:
+                stored_type_name = self.metadata[key]['data_type']
+                if stored_type_name == 'DataFrame':
+                    extension = '.csv'
+                elif stored_type_name == 'dict':
+                    extension = '.json'
+                elif stored_type_name == 'str':
+                    extension = '.txt'
+                else:
+                    # Try to find extension from registered handlers
+                    extension = '.dat'  # Default extension
+                    for registered_type, handler in self.handlers.items():
+                        if registered_type.__name__ == stored_type_name:
+                            extension = handler.get_file_extension()
+                            break
+            else:
+                extension = '.dat'  # Default extension
+        else:
+            try:
+                handler = self._get_handler(data_type)
+                extension = handler.get_file_extension()
+            except ValueError:
+                extension = '.dat'  # Default extension
+
+        return self.cache_dir / f"{key}{extension}"
 
     def _is_expired(self, key: str) -> bool:
         """Check if data is expired"""
@@ -55,32 +113,27 @@ class DataFrameCache:
         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:
+    def save(self, key: str, data: Union[pd.DataFrame, dict, Any], expire_hours: Optional[float] = None,
+             **handler_kwargs) -> bool:
         """
-        Save DataFrame to cache
-        
+        Save data to cache
+
         Args:
             key: Cache key name
-            df: DataFrame to save
+            data: Data to save (DataFrame, dict, or other supported types)
             expire_hours: Expiration time in hours, None means never expires
-            **csv_kwargs: Additional parameters passed to pandas to_csv
-            
+            **handler_kwargs: Additional parameters passed to the data handler
+
         Returns:
             bool: Whether save was successful
         """
         try:
-            file_path = self._get_file_path(key)
+            data_type = type(data)
+            handler = self._get_handler(data_type)
+            file_path = self._get_file_path(key, data_type)
 
-            # 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)
+            # Save data using appropriate handler
+            handler_metadata = handler.save(data, file_path, **handler_kwargs)
 
             # Update metadata
             current_time = datetime.now()
@@ -88,29 +141,28 @@ class DataFrameCache:
                 '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)
+                'data_type': data_type.__name__,
+                **handler_metadata
             }
 
             self._save_metadata()
             return True
 
         except Exception as e:
-            print(f"Failed to save DataFrame: {e}")
+            print(f"Failed to save data: {e}")
             return False
 
-    def load(self, key: str, auto_clean_expired: bool = True, **csv_kwargs) -> Optional[pd.DataFrame]:
+    def load(self, key: str, auto_clean_expired: bool = True, **handler_kwargs) -> Optional[Any]:
         """
-        Load DataFrame from cache
-        
+        Load data 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
-            
+            **handler_kwargs: Additional parameters passed to the data handler
+
         Returns:
-            Optional[pd.DataFrame]: Loaded DataFrame, returns None if not exists or expired
+            Optional[Any]: Loaded data, returns None if not exists or expired
         """
         try:
             # Check if expired
@@ -124,34 +176,56 @@ class DataFrameCache:
             if not file_path.exists():
                 return None
 
-            # Set default CSV parameters
-            csv_params = {
-                'encoding': 'utf-8'
-            }
-            csv_params.update(csv_kwargs)
+            # Get data type from metadata
+            if key not in self.metadata or 'data_type' not in self.metadata[key]:
+                print(f"No data type information found for key '{key}'")
+                return None
+
+            data_type_name = self.metadata[key]['data_type']
+
+            # Map type name back to actual type
+            if data_type_name == 'DataFrame':
+                data_type = pd.DataFrame
+            elif data_type_name == 'dict':
+                data_type = dict
+            elif data_type_name == 'str':
+                data_type = str
+            else:
+                # For other custom types, try to find a handler by checking registered types
+                data_type = None
+                for registered_type in self.handlers.keys():
+                    if registered_type.__name__ == data_type_name:
+                        data_type = registered_type
+                        break
+
+                if data_type is None:
+                    print(f"Unknown data type: {data_type_name}")
+                    return None
 
-            # Read CSV file
-            df = pd.read_csv(file_path, **csv_params)
+            handler = self._get_handler(data_type)
+
+            # Load data using appropriate handler
+            data = handler.load(file_path, **handler_kwargs)
 
             # Update last access time
             if key in self.metadata:
                 self.metadata[key]['last_accessed'] = datetime.now().isoformat()
                 self._save_metadata()
 
-            return df
+            return data
 
         except Exception as e:
-            print(f"Failed to load DataFrame: {e}")
+            print(f"Failed to load data: {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
         """
@@ -164,17 +238,17 @@ class DataFrameCache:
     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
+            # Delete data file
             if file_path.exists():
                 file_path.unlink()
 
@@ -192,7 +266,7 @@ class DataFrameCache:
     def clean_expired(self) -> int:
         """
         Clean all expired caches
-        
+
         Returns:
             int: Number of cleaned caches
         """
@@ -212,10 +286,10 @@ class DataFrameCache:
     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.
         """
@@ -232,10 +306,10 @@ class DataFrameCache:
     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
         """
@@ -254,7 +328,7 @@ class DataFrameCache:
     def get_cache_stats(self) -> Dict[str, Any]:
         """
         Get cache statistics
-        
+
         Returns:
             Dict: Cache statistics information
         """
@@ -280,14 +354,15 @@ class DataFrameCache:
     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()
+            # Delete all data files (CSV, JSON, and other supported formats)
+            for data_file in self.cache_dir.glob("*"):
+                if data_file.is_file() and data_file.name != "metadata.json":
+                    data_file.unlink()
 
             # Clear metadata
             self.metadata = {}
@@ -298,34 +373,3 @@ class DataFrameCache:
         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/storage/vector_store/__init__.py

@@ -1,3 +1,3 @@
-from flowllm.storage.vector_store.chroma_vector_store import ChromaVectorStore
-from flowllm.storage.vector_store.es_vector_store import EsVectorStore
-from flowllm.storage.vector_store.local_vector_store import LocalVectorStore
+from .chroma_vector_store import ChromaVectorStore
+from .es_vector_store import EsVectorStore
+from .local_vector_store import LocalVectorStore

flowllm/storage/vector_store/es_vector_store.py

@@ -13,7 +13,7 @@ from flowllm.storage.vector_store.local_vector_store import LocalVectorStore
 
 @C.register_vector_store("elasticsearch")
 class EsVectorStore(LocalVectorStore):
-    hosts: str | List[str] = Field(default_factory=lambda: os.getenv("ES_HOSTS", "http://localhost:9200"))
+    hosts: str | List[str] = Field(default_factory=lambda: os.getenv("FLOW_ES_HOSTS", "http://localhost:9200"))
     basic_auth: str | Tuple[str, str] | None = Field(default=None)
     retrieve_filters: List[dict] = []
     _client: Elasticsearch = PrivateAttr()
@@ -160,7 +160,6 @@ class EsVectorStore(LocalVectorStore):
         if refresh:
             self.refresh(workspace_id=workspace_id)
 
-
 def main():
     from flowllm.utils.common_utils import load_env
     from flowllm.embedding_model import OpenAICompatibleEmbeddingModel

flowllm/storage/vector_store/local_vector_store.py

@@ -203,7 +203,6 @@ class LocalVectorStore(BaseVectorStore):
         self._dump_to_path(nodes=all_nodes, workspace_id=workspace_id, path=self.store_path, **kwargs)
         logger.info(f"delete workspace_id={workspace_id} before_size={before_size} after_size={after_size}")
 
-
 def main():
     from flowllm.utils.common_utils import load_env
     from flowllm.embedding_model import OpenAICompatibleEmbeddingModel

flowllm/utils/common_utils.py

@@ -40,25 +40,13 @@ def load_env(path: str | Path = 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)
+    else:
+        for i in range(5):
+            path = Path("../" * i + ".env")
+            if path.exists():
+                logger.info(f"using path={path}")
+                _load_env(path)
+                return
+
+        raise FileNotFoundError(".env not found")

flowllm/utils/fetch_url.py

@@ -1,12 +1,14 @@
-import requests
-from bs4 import BeautifulSoup
-import urllib3
 import random
 import time
 
+import requests
+import urllib3
+from bs4 import BeautifulSoup
+
 # 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 = [
@@ -21,7 +23,7 @@ def get_random_headers():
         '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',
@@ -29,13 +31,13 @@ def get_random_headers():
         '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',
@@ -44,7 +46,7 @@ def get_random_headers():
         'Connection': 'keep-alive',
         'Upgrade-Insecure-Requests': '1',
     }
-    
+
     # Randomly add some optional headers
     if random.random() > 0.5:
         headers['DNT'] = '1'
@@ -54,9 +56,10 @@ def get_random_headers():
         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
@@ -72,12 +75,12 @@ def fetch_webpage_text(url, min_delay=1, max_delay=3):
     # 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)
@@ -106,8 +109,9 @@ def fetch_webpage_text(url, min_delay=1, max_delay=3):
     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)
+    print(text)

flowllm/utils/llm_utils.py

@@ -0,0 +1,28 @@
+from typing import List
+
+from flowllm.enumeration.role import Role
+from flowllm.schema.message import Message
+
+
+def merge_messages_content(messages: List[Message | dict]) -> str:
+    content_collector = []
+    for i, message in enumerate(messages):
+        if isinstance(message, dict):
+            message = Message(**message)
+
+        if message.role is Role.ASSISTANT:
+            line = f"### step.{i} role={message.role.value} content=\n{message.reasoning_content}\n\n{message.content}\n"
+            if message.tool_calls:
+                for tool_call in message.tool_calls:
+                    line += f" - tool call={tool_call.name}\n   params={tool_call.arguments}\n"
+            content_collector.append(line)
+
+        elif message.role is Role.USER:
+            line = f"### step.{i} role={message.role.value} content=\n{message.content}\n"
+            content_collector.append(line)
+
+        elif message.role is Role.TOOL:
+            line = f"### step.{i} role={message.role.value} tool call result=\n{message.content}\n"
+            content_collector.append(line)
+
+    return "\n".join(content_collector)

flowllm/utils/ridge_v2.py

@@ -0,0 +1,54 @@
+import numpy as np
+from sklearn.linear_model import Ridge
+from sklearn.preprocessing import StandardScaler
+
+
+class RidgeV2:
+
+    def __init__(self, bound: int = 4.6, use_ridge_v2=True, **kwargs):
+        self.bound: int = bound
+        self.use_ridge_v2: bool = use_ridge_v2
+        self.kwargs: dict = kwargs
+        self.model = Ridge(**self.kwargs)
+
+    def clear(self):
+        self.model = Ridge(**self.kwargs)
+        return self
+
+    def sigmoid(self, x):
+        x = np.asarray(x, dtype=float)
+        x = np.clip(x, -self.bound, self.bound)
+        return 1 / (1 + np.exp(-x))
+
+    def inv_sigmoid(self, p):
+        p = np.asarray(p, dtype=float)
+        p = np.clip(p, self.sigmoid(-self.bound), self.sigmoid(self.bound))
+        return np.log(p / (1 - p))
+
+    def fit(self, x, y, sample_weight=None):
+        if self.use_ridge_v2:
+            return self.model.fit(x, self.inv_sigmoid(y), sample_weight=sample_weight)
+        else:
+            return self.model.fit(x, y, sample_weight=sample_weight)
+
+    def predict(self, x):
+        if self.use_ridge_v2:
+            return self.sigmoid(self.model.predict(x))
+        else:
+            return self.model.predict(x)
+
+    def fit_and_predict(self,
+                        train_x_nd: np.ndarray,
+                        train_y_nd: np.ndarray,
+                        test_x_nd: np.ndarray,
+                        check_y: bool = True):
+        if check_y:
+            assert np.all((train_y_nd >= 0) & (train_y_nd <= 1))
+
+        scaler = StandardScaler()
+        scaler.fit(train_x_nd)
+        train_x_nd = scaler.transform(train_x_nd)
+        test_x_nd = scaler.transform(test_x_nd)
+        self.model.fit(train_x_nd, train_y_nd)
+        pred_y_nd = self.model.predict(test_x_nd)
+        return np.minimum(np.maximum(pred_y_nd, 0), 1)