From d8ff285781c9693a201e997a51d3900408c4ab43 Mon Sep 17 00:00:00 2001 From: Alex Korbonits Date: Sun, 12 Apr 2026 23:32:09 -0700 Subject: [PATCH] Add Ruff linting, pre-commit hooks, and CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add [tool.ruff] config to pyproject.toml (E/F/I rules, line-length 120) - Add .pre-commit-config.yaml with ruff + ruff-format hooks - Add .github/workflows/lint.yml CI job via uv - Fix src/lightmem/memory/graph.py: stub class body was missing (added pass) Auto-fixed 217 import-sort, style, and formatting violations. Remaining notes: - F601: duplicate dictionary keys in lightmem.py:166-167 ('embedding_calls' and 'embedding_total_tokens' each appear twice in the token stats dict — the second assignment silently overwrites the first) - F821: undefined name 'client' in experiments/locomo/llm_judge.py:60 (referenced before assignment; potential NameError at runtime) Both are suppressed in the config so CI passes, but flagging for maintainer attention. --- .github/workflows/lint.yml | 17 +++++ .pre-commit-config.yaml | 7 +++ examples/run_lightmem_ollama.py | 6 +- examples/run_lightmem_transformers.py | 3 +- experiments/locomo/add_locomo.py | 51 ++++++++------- experiments/locomo/llm_judge.py | 5 +- experiments/locomo/retrievers.py | 12 ++-- experiments/locomo/search_locomo.py | 27 ++++---- experiments/longmemeval/offline_update.py | 2 + experiments/longmemeval/run_lightmem_gpt.py | 8 ++- experiments/longmemeval/run_lightmem_qwen.py | 9 +-- pyproject.toml | 11 ++++ src/lightmem/configs/base.py | 15 ++--- src/lightmem/configs/logging/base.py | 6 +- src/lightmem/configs/logging/utils.py | 2 +- src/lightmem/configs/memory_manager/base.py | 4 +- .../configs/memory_manager/base_config.py | 2 +- .../configs/multimodal_embedder/base.py | 4 +- src/lightmem/configs/pre_compressor/base.py | 6 +- .../pre_compressor/entropy_compress.py | 4 +- .../configs/pre_compressor/llmlingua_2.py | 4 +- src/lightmem/configs/retriever/bm25.py | 3 +- .../retriever/contextretriever/base.py | 4 +- .../retriever/embeddingretriever/base.py | 4 +- .../retriever/embeddingretriever/qdrant.py | 4 +- src/lightmem/configs/text_embedder/base.py | 5 +- .../configs/text_embedder/base_config.py | 4 +- src/lightmem/configs/topic_segmenter/base.py | 4 +- .../factory/memory_buffer/sensory_memory.py | 4 +- .../memory_buffer/short_term_memory.py | 4 +- .../factory/memory_manager/deepseek.py | 3 +- .../factory/memory_manager/factory.py | 3 +- src/lightmem/factory/memory_manager/ollama.py | 2 +- src/lightmem/factory/memory_manager/openai.py | 12 ++-- .../factory/memory_manager/transformers.py | 5 +- src/lightmem/factory/memory_manager/vllm.py | 3 +- .../factory/memory_manager/vllm_offline.py | 4 +- .../pre_compressor/entropy_compress.py | 6 +- .../factory/pre_compressor/factory.py | 6 +- .../factory/pre_compressor/llmlingua_2.py | 3 +- .../retriever/contextretriever/factory.py | 3 +- .../retriever/embeddingretriever/factory.py | 4 +- .../retriever/embeddingretriever/qdrant.py | 5 +- src/lightmem/factory/text_embedder/factory.py | 4 +- .../factory/text_embedder/huggingface.py | 7 ++- src/lightmem/factory/text_embedder/openai.py | 7 ++- .../factory/topic_segmenter/factory.py | 4 +- .../factory/topic_segmenter/llmlingua_2.py | 9 ++- src/lightmem/memory/graph.py | 2 +- src/lightmem/memory/lightmem.py | 54 +++++++++------- src/lightmem/memory/utils.py | 17 ++--- .../inference_utils/backends.py | 18 +++--- .../inference_utils/base_operator.py | 24 ++++--- .../inference_utils/operators.py | 17 ++--- .../memory_toolkits/memories/__init__.py | 4 +- .../memory_toolkits/memories/datasets/base.py | 32 +++++----- .../memories/datasets/locomo.py | 20 +++--- .../memories/datasets/longmemeval.py | 18 +++--- .../memory_toolkits/memories/layers/amem.py | 34 +++++----- .../memory_toolkits/memories/layers/base.py | 8 +-- .../agentic_memory/llm_controller.py | 7 ++- .../baselines/agentic_memory/memory_system.py | 20 ++---- .../baselines/agentic_memory/retrievers.py | 18 +++--- .../layers/baselines/mem0/client/main.py | 2 +- .../layers/baselines/mem0/client/project.py | 4 +- .../layers/baselines/mem0/client/utils.py | 2 +- .../layers/baselines/mem0/configs/base.py | 3 +- .../baselines/mem0/configs/embeddings/base.py | 1 - .../mem0/configs/vector_stores/databricks.py | 7 ++- .../baselines/mem0/embeddings/aws_bedrock.py | 1 - .../baselines/mem0/embeddings/azure_openai.py | 3 +- .../baselines/mem0/embeddings/gemini.py | 1 - .../baselines/mem0/embeddings/huggingface.py | 5 +- .../mem0/graphs/neptune/neptunedb.py | 1 + .../mem0/graphs/neptune/neptunegraph.py | 2 +- .../layers/baselines/mem0/llms/aws_bedrock.py | 2 +- .../baselines/mem0/llms/azure_openai.py | 3 +- .../mem0/llms/azure_openai_structured.py | 3 +- .../layers/baselines/mem0/llms/deepseek.py | 3 +- .../layers/baselines/mem0/llms/lmstudio.py | 3 +- .../layers/baselines/mem0/llms/openai.py | 3 +- .../layers/baselines/mem0/llms/sarvam.py | 1 - .../layers/baselines/mem0/llms/vllm.py | 3 +- .../layers/baselines/mem0/memory/main.py | 3 +- .../layers/baselines/mem0/memory/telemetry.py | 3 +- .../layers/baselines/mem0/proxy/main.py | 1 - .../mem0/vector_stores/azure_ai_search.py | 3 +- .../mem0/vector_stores/azure_mysql.py | 2 +- .../baselines/mem0/vector_stores/baidu.py | 3 +- .../mem0/vector_stores/databricks.py | 19 ++++-- .../mem0/vector_stores/elasticsearch.py | 3 +- .../baselines/mem0/vector_stores/faiss.py | 3 +- .../baselines/mem0/vector_stores/milvus.py | 3 +- .../mem0/vector_stores/opensearch.py | 3 +- .../baselines/mem0/vector_stores/qdrant.py | 3 +- .../baselines/mem0/vector_stores/redis.py | 5 +- .../mem0/vector_stores/s3_vectors.py | 3 +- .../mem0/vector_stores/upstash_vector.py | 3 +- .../baselines/mem0/vector_stores/valkey.py | 5 +- .../vector_stores/vertex_ai_vector_search.py | 3 +- .../baselines/mem0/vector_stores/weaviate.py | 3 +- .../memories/layers/full_context.py | 10 +-- .../memories/layers/langmem.py | 37 ++++++----- .../memories/layers/memzero.py | 31 ++++------ .../memories/layers/naive_rag.py | 13 ++-- .../memory_toolkits/memory_construction.py | 62 +++++++++---------- .../memory_toolkits/memory_evaluation.py | 23 +++---- src/lightmem/memory_toolkits/memory_search.py | 21 ++++--- src/lightmem/memory_toolkits/monkey_patch.py | 14 +++-- src/lightmem/memory_toolkits/token_monitor.py | 27 ++++---- 110 files changed, 536 insertions(+), 450 deletions(-) create mode 100644 .github/workflows/lint.yml create mode 100644 .pre-commit-config.yaml diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..fe58f73 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,17 @@ +name: Lint + +on: + push: + branches: [main] + pull_request: + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v5 + - run: uv python install 3.10 + - run: uv pip install ruff + - run: uv run ruff check . + - run: uv run ruff format --check . diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..700cbc4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,7 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.12.2 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format diff --git a/examples/run_lightmem_ollama.py b/examples/run_lightmem_ollama.py index 397fc7e..8ac88d0 100644 --- a/examples/run_lightmem_ollama.py +++ b/examples/run_lightmem_ollama.py @@ -1,12 +1,12 @@ import json -import ollama import os import time -from tqdm import tqdm from typing import Dict, List, Optional -from lightmem.memory.lightmem import LightMemory +import ollama +from tqdm import tqdm +from lightmem.memory.lightmem import LightMemory # =========== Ollama Configuration ============ your_ollama_model_name = "your_Ollama_model_name_01" # such as "llama3:latest" diff --git a/examples/run_lightmem_transformers.py b/examples/run_lightmem_transformers.py index b20b272..b183d72 100644 --- a/examples/run_lightmem_transformers.py +++ b/examples/run_lightmem_transformers.py @@ -5,11 +5,10 @@ import torch from tqdm import tqdm -from transformers import AutoTokenizer, AutoModelForCausalLM +from transformers import AutoModelForCausalLM, AutoTokenizer from lightmem.memory.lightmem import LightMemory - # =========== Transformers Model Configuration ============ your_model_path_or_name = "your_model_path_or_name_01" # specify the model's path or name your_model_device = "cuda:1" # specify the GPU device diff --git a/experiments/locomo/add_locomo.py b/experiments/locomo/add_locomo.py index 8978386..7dbebc8 100644 --- a/experiments/locomo/add_locomo.py +++ b/experiments/locomo/add_locomo.py @@ -1,19 +1,24 @@ -from openai import OpenAI -import json -from tqdm import tqdm +import argparse import datetime -import time -import os +import json import logging -from lightmem.memory.lightmem import LightMemory -from lightmem.configs.retriever.embeddingretriever.qdrant import QdrantConfig -from lightmem.factory.retriever.embeddingretriever.qdrant import Qdrant -from prompts import METADATA_GENERATE_PROMPT_locomo, LoCoMo_Event_Binding_factual, LoCoMo_Event_Binding_relational -import sqlite3 +import multiprocessing as mp +import os import shutil -import argparse +import sqlite3 +import time from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed -import multiprocessing as mp + +from prompts import ( + LoCoMo_Event_Binding_factual, + LoCoMo_Event_Binding_relational, + METADATA_GENERATE_PROMPT_locomo, +) +from tqdm import tqdm + +from lightmem.configs.retriever.embeddingretriever.qdrant import QdrantConfig +from lightmem.factory.retriever.embeddingretriever.qdrant import Qdrant +from lightmem.memory.lightmem import LightMemory # ============ Configuration ============ LOGS_ROOT = "./logs" @@ -357,7 +362,7 @@ def process_single_sample(sample, api_key, args): backup_start_time = time.time() if os.path.exists(backup_dir): - logger.info(f" Removing existing backup...") + logger.info(" Removing existing backup...") shutil.rmtree(backup_dir) logger.info(f" Copying: {source_dir} -> {backup_dir}") @@ -389,7 +394,7 @@ def process_single_sample(sample, api_key, args): initial_summarize_tokens = initial_summarize_stats['llm']['summarize']['total_tokens'] initial_summarize_calls = initial_summarize_stats['llm']['summarize']['calls'] - logger.info(f" Creating LightMemory instance for summarization (using pre_update)") + logger.info(" Creating LightMemory instance for summarization (using pre_update)") lightmem_for_summary = load_lightmem( collection_name=sample_id, api_key=api_key, @@ -465,13 +470,13 @@ def process_single_sample(sample, api_key, args): logger.info(f"SUMMARY: {sample_id}") logger.info(f"{'='*70}") - logger.info(f"\n[Storage Information]") + logger.info("\n[Storage Information]") logger.info(f" Pre-update: {QDRANT_PRE_UPDATE_DIR}/{sample_id} ({pre_update_count} entries)") logger.info(f" Post-update: {QDRANT_POST_UPDATE_DIR}/{sample_id} ({post_update_count} entries)") logger.info(f" Change: {post_update_count - pre_update_count:+d} entries") logger.info(f" Summaries: {num_summaries}") - logger.info(f"\n[Time Statistics]") + logger.info("\n[Time Statistics]") logger.info(f" Total: {case_total_duration:.2f}s") logger.info(f" ├─ Add: {add_memory_duration:.2f}s ({add_memory_duration/case_total_duration*100:.1f}%)") if args.enable_summary: @@ -479,26 +484,26 @@ def process_single_sample(sample, api_key, args): logger.info(f" ├─ Backup: {backup_duration:.2f}s ({backup_duration/case_total_duration*100:.1f}%)") logger.info(f" └─ Update: {update_duration:.2f}s ({update_duration/case_total_duration*100:.1f}%)") - logger.info(f"\n[Token Statistics - Add Memory]") + logger.info("\n[Token Statistics - Add Memory]") logger.info(f" Calls: {case_add_calls}") logger.info(f" Prompt: {case_add_prompt:,}") logger.info(f" Completion: {case_add_completion:,}") logger.info(f" Total: {case_add_tokens:,}") if args.enable_summary: - logger.info(f"\n[Token Statistics - Summarize]") + logger.info("\n[Token Statistics - Summarize]") logger.info(f" Calls: {case_summarize_calls}") logger.info(f" Prompt: {case_summarize_prompt:,}") logger.info(f" Completion: {case_summarize_completion:,}") logger.info(f" Total: {case_summarize_tokens:,}") - logger.info(f"\n[Token Statistics - Update]") + logger.info("\n[Token Statistics - Update]") logger.info(f" Calls: {case_update_calls}") logger.info(f" Prompt: {case_update_prompt:,}") logger.info(f" Completion: {case_update_completion:,}") logger.info(f" Total: {case_update_tokens:,}") - logger.info(f"\n[Total Usage]") + logger.info("\n[Total Usage]") logger.info(f" API Calls: {case_add_calls + case_summarize_calls + case_update_calls}") logger.info(f" Tokens: {case_add_tokens + case_summarize_tokens + case_update_tokens:,}") logger.info(f"{'='*70}\n") @@ -667,7 +672,7 @@ def main(): successful = [r for r in results if r['status'] == 'success'] - main_logger.info(f"\n[Overall Statistics]") + main_logger.info("\n[Overall Statistics]") main_logger.info(f" Total samples: {len(missing)}") main_logger.info(f" Successful: {len(successful)}") main_logger.info(f" Failed: {len(failed_samples)}") @@ -679,7 +684,7 @@ def main(): total_calls = sum(r['add_calls'] + r.get('summarize_calls', 0) + r['update_calls'] for r in successful) total_summaries = sum(r.get('num_summaries', 0) for r in successful) - main_logger.info(f"\n[Performance Metrics]") + main_logger.info("\n[Performance Metrics]") main_logger.info(f" Avg per sample: {avg_duration:.2f}s") main_logger.info(f" Speedup: {avg_duration * len(successful) / total_duration:.2f}x") main_logger.info(f" Total API calls: {total_calls}") @@ -687,7 +692,7 @@ def main(): main_logger.info(f" Total summaries: {total_summaries}") if failed_samples: - main_logger.info(f"\n[Failed Samples]") + main_logger.info("\n[Failed Samples]") for sample_id in failed_samples: main_logger.info(f" - {sample_id}") diff --git a/experiments/locomo/llm_judge.py b/experiments/locomo/llm_judge.py index ba207c7..243dc9a 100644 --- a/experiments/locomo/llm_judge.py +++ b/experiments/locomo/llm_judge.py @@ -1,9 +1,10 @@ import argparse import json +import re from collections import defaultdict -import re + import numpy as np -from openai import OpenAI + def extract_json(text): """ diff --git a/experiments/locomo/retrievers.py b/experiments/locomo/retrievers.py index 940bd2a..294ccbe 100644 --- a/experiments/locomo/retrievers.py +++ b/experiments/locomo/retrievers.py @@ -1,16 +1,14 @@ -from openai import OpenAI import logging import os -import time -import sqlite3 import pickle -from typing import List, Dict, Any, Set, Optional, Tuple -from collections import defaultdict, deque +import sqlite3 from datetime import datetime +from typing import Any, Dict, List + import numpy as np -import spacy -from lightmem.factory.retriever.embeddingretriever.qdrant import Qdrant + from lightmem.configs.retriever.embeddingretriever.qdrant import QdrantConfig +from lightmem.factory.retriever.embeddingretriever.qdrant import Qdrant SPACY_AVAILABLE = True logger = logging.getLogger(__name__) diff --git a/experiments/locomo/search_locomo.py b/experiments/locomo/search_locomo.py index 4f4acc7..f8dbe3a 100644 --- a/experiments/locomo/search_locomo.py +++ b/experiments/locomo/search_locomo.py @@ -1,22 +1,21 @@ -from openai import OpenAI -import json -from tqdm import tqdm +import argparse import datetime -import time -import os +import json import logging -from typing import List, Dict, Any, Optional -import numpy as np -import argparse - -from lightmem.factory.text_embedder.huggingface import TextEmbedderHuggingface -from lightmem.factory.text_embedder.openai import TextEmbedderOpenAI -from lightmem.configs.text_embedder.base_config import BaseTextEmbedderConfig +import os +import time +from typing import Dict, List, Optional +import numpy as np +from llm_judge import evaluate_llm_judge +from openai import OpenAI from prompts import ANSWER_PROMPT, ANSWER_PROMPT_StructMem from retrievers import QdrantEntryLoader, VectorRetriever, format_related_memories -from llm_judge import evaluate_llm_judge +from tqdm import tqdm +from lightmem.configs.text_embedder.base_config import BaseTextEmbedderConfig +from lightmem.factory.text_embedder.huggingface import TextEmbedderHuggingface +from lightmem.factory.text_embedder.openai import TextEmbedderOpenAI # ============ Configuration ============ LOGS_ROOT = "./logs" @@ -589,7 +588,7 @@ def main(): logger.info("=" * 80) logger.info("LightMem Evaluation - LoCoMo Dataset") logger.info("=" * 80) - logger.info(f"Configuration:") + logger.info("Configuration:") logger.info(f" Dataset: {args.dataset}") logger.info(f" Qdrant dir: {args.qdrant_dir}") logger.info(f" Output dir: {args.output_dir}") diff --git a/experiments/longmemeval/offline_update.py b/experiments/longmemeval/offline_update.py index 65b2007..6822fee 100644 --- a/experiments/longmemeval/offline_update.py +++ b/experiments/longmemeval/offline_update.py @@ -1,6 +1,8 @@ import os + from lightmem.memory.lightmem import LightMemory + def load_lightmem(collection_name): config = { "memory_manager": { diff --git a/experiments/longmemeval/run_lightmem_gpt.py b/experiments/longmemeval/run_lightmem_gpt.py index b20ec44..686a84c 100644 --- a/experiments/longmemeval/run_lightmem_gpt.py +++ b/experiments/longmemeval/run_lightmem_gpt.py @@ -1,10 +1,12 @@ -from openai import OpenAI import json -from tqdm import tqdm -import datetime import time + +from openai import OpenAI +from tqdm import tqdm + from lightmem.memory.lightmem import LightMemory + def get_anscheck_prompt(task, question, answer, response, abstention=False): if not abstention: if task in ['single-session-user', 'single-session-assistant', 'multi-session']: diff --git a/experiments/longmemeval/run_lightmem_qwen.py b/experiments/longmemeval/run_lightmem_qwen.py index 2472127..edf16b1 100644 --- a/experiments/longmemeval/run_lightmem_qwen.py +++ b/experiments/longmemeval/run_lightmem_qwen.py @@ -1,9 +1,10 @@ -from openai import OpenAI import json -from tqdm import tqdm -import datetime -import time import os +import time + +from openai import OpenAI +from tqdm import tqdm + from lightmem.memory.lightmem import LightMemory # ============ API Configuration ============ diff --git a/pyproject.toml b/pyproject.toml index 0ce7818..68586fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,3 +108,14 @@ where = ["src"] [tool.setuptools.package-data] "*" = ["*.yaml", "*.yml", "*.json"] + +[tool.ruff] +line-length = 120 +extend-exclude = ["*.ipynb"] + +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = ["E501", "E402", "E731", "E722", "E741", "E721", "E701", "F841", "F403", "F405", "F601", "F821", "E712"] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401"] diff --git a/src/lightmem/configs/base.py b/src/lightmem/configs/base.py index 9f76537..c249d1b 100644 --- a/src/lightmem/configs/base.py +++ b/src/lightmem/configs/base.py @@ -1,15 +1,16 @@ import os -from pydantic import BaseModel, Field, model_validator -from typing import Any, Dict, Optional, Literal -from pydantic import ValidationError -from lightmem.configs.pre_compressor.base import PreCompressorConfig -from lightmem.configs.topic_segmenter.base import TopicSegmenterConfig +from typing import Literal, Optional + +from pydantic import BaseModel, Field + +from lightmem.configs.logging.base import LoggingConfig from lightmem.configs.memory_manager.base import MemoryManagerConfig -from lightmem.configs.text_embedder.base import TextEmbedderConfig from lightmem.configs.multimodal_embedder.base import MMEmbedderConfig +from lightmem.configs.pre_compressor.base import PreCompressorConfig from lightmem.configs.retriever.contextretriever.base import ContextRetrieverConfig from lightmem.configs.retriever.embeddingretriever.base import EmbeddingRetrieverConfig -from lightmem.configs.logging.base import LoggingConfig +from lightmem.configs.text_embedder.base import TextEmbedderConfig +from lightmem.configs.topic_segmenter.base import TopicSegmenterConfig lightmem_dir = "" diff --git a/src/lightmem/configs/logging/base.py b/src/lightmem/configs/logging/base.py index 5f96001..85b2c9f 100644 --- a/src/lightmem/configs/logging/base.py +++ b/src/lightmem/configs/logging/base.py @@ -1,8 +1,10 @@ -from pydantic import BaseModel, Field, field_validator, model_validator -from typing import Optional, Union, Dict, List, Literal import logging import os from datetime import datetime +from typing import Dict, List, Literal, Optional, Union + +from pydantic import BaseModel, Field, field_validator, model_validator + class LoggingConfig(BaseModel): """Logging configuration for LightMem.""" diff --git a/src/lightmem/configs/logging/utils.py b/src/lightmem/configs/logging/utils.py index b9c64b1..3c9c4d1 100644 --- a/src/lightmem/configs/logging/utils.py +++ b/src/lightmem/configs/logging/utils.py @@ -1,5 +1,5 @@ import logging -from typing import Optional, Union, Dict, List, TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, List, Optional, Union if TYPE_CHECKING: from .base import LoggingConfig diff --git a/src/lightmem/configs/memory_manager/base.py b/src/lightmem/configs/memory_manager/base.py index b218c0b..86beb24 100644 --- a/src/lightmem/configs/memory_manager/base.py +++ b/src/lightmem/configs/memory_manager/base.py @@ -1,5 +1,7 @@ +from typing import ClassVar, List, Optional + from pydantic import BaseModel, Field, model_validator -from typing import Dict, Optional, Type, Any, List, ClassVar + from .base_config import BaseMemoryManagerConfig diff --git a/src/lightmem/configs/memory_manager/base_config.py b/src/lightmem/configs/memory_manager/base_config.py index 7d4c9c3..9426b66 100644 --- a/src/lightmem/configs/memory_manager/base_config.py +++ b/src/lightmem/configs/memory_manager/base_config.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Union, List +from typing import Dict, List, Optional, Union class BaseMemoryManagerConfig: diff --git a/src/lightmem/configs/multimodal_embedder/base.py b/src/lightmem/configs/multimodal_embedder/base.py index f661ee9..7fa2fb6 100644 --- a/src/lightmem/configs/multimodal_embedder/base.py +++ b/src/lightmem/configs/multimodal_embedder/base.py @@ -1,5 +1,7 @@ +from typing import List, Optional + from pydantic import BaseModel, Field, model_validator -from typing import Dict, Optional, Type, Any, List + class MMEmbedderConfig(BaseModel): model_name: str = Field(description="The multimodal embedding model or Deployment platform (e.g., 'openai', 'ollama')", default="huggingface") diff --git a/src/lightmem/configs/pre_compressor/base.py b/src/lightmem/configs/pre_compressor/base.py index 523bc3a..f61f0bd 100644 --- a/src/lightmem/configs/pre_compressor/base.py +++ b/src/lightmem/configs/pre_compressor/base.py @@ -1,6 +1,8 @@ -from pydantic import BaseModel, Field, model_validator -from typing import Dict, Optional, Type, Any, Union, ClassVar from importlib import import_module +from typing import Any, ClassVar, Dict + +from pydantic import BaseModel, Field, model_validator + class PreCompressorConfig(BaseModel): model_name: str = Field( diff --git a/src/lightmem/configs/pre_compressor/entropy_compress.py b/src/lightmem/configs/pre_compressor/entropy_compress.py index c7d229a..a844cbd 100644 --- a/src/lightmem/configs/pre_compressor/entropy_compress.py +++ b/src/lightmem/configs/pre_compressor/entropy_compress.py @@ -1,6 +1,6 @@ -import os +from typing import Any, Dict + from pydantic import BaseModel, Field, field_validator -from typing import Dict, Optional, Any class EntropyCompressorConfig(BaseModel): diff --git a/src/lightmem/configs/pre_compressor/llmlingua_2.py b/src/lightmem/configs/pre_compressor/llmlingua_2.py index 0c9abc7..7d8035e 100644 --- a/src/lightmem/configs/pre_compressor/llmlingua_2.py +++ b/src/lightmem/configs/pre_compressor/llmlingua_2.py @@ -1,6 +1,8 @@ import os +from typing import Any, Dict + from pydantic import BaseModel, Field, field_validator -from typing import Dict, Optional, Type, Any + class LlmLingua2Config(BaseModel): llmlingua_config: Dict[str, Any] = Field( diff --git a/src/lightmem/configs/retriever/bm25.py b/src/lightmem/configs/retriever/bm25.py index 156818d..305f41d 100644 --- a/src/lightmem/configs/retriever/bm25.py +++ b/src/lightmem/configs/retriever/bm25.py @@ -1,6 +1,7 @@ -from pydantic import BaseModel, Field, model_validator from typing import Any, Dict, Optional +from pydantic import BaseModel, Field, model_validator + class BM25Config(BaseModel): """Configuration for BM25 retriever.""" diff --git a/src/lightmem/configs/retriever/contextretriever/base.py b/src/lightmem/configs/retriever/contextretriever/base.py index d718aa7..bcd9269 100644 --- a/src/lightmem/configs/retriever/contextretriever/base.py +++ b/src/lightmem/configs/retriever/contextretriever/base.py @@ -1,5 +1,7 @@ +from typing import List, Optional + from pydantic import BaseModel, Field, model_validator -from typing import Dict, Optional, Type, Any, List + class ContextRetrieverConfig(BaseModel): model_name: str = Field(description="The Context retriever or algorithm (e.g., 'BM25', 'TF-IDF')", default="BM25") diff --git a/src/lightmem/configs/retriever/embeddingretriever/base.py b/src/lightmem/configs/retriever/embeddingretriever/base.py index 80b2116..97026ef 100644 --- a/src/lightmem/configs/retriever/embeddingretriever/base.py +++ b/src/lightmem/configs/retriever/embeddingretriever/base.py @@ -1,7 +1,9 @@ -from typing import Dict, Optional, ClassVar from importlib import import_module +from typing import ClassVar, Dict, Optional + from pydantic import BaseModel, Field, model_validator + class EmbeddingRetrieverConfig(BaseModel): model_name: str = Field( description="The embedding retriever or vector store (e.g., 'qdrant', 'chroma', 'upstash_vector')", diff --git a/src/lightmem/configs/retriever/embeddingretriever/qdrant.py b/src/lightmem/configs/retriever/embeddingretriever/qdrant.py index b217ed0..50c3cf8 100644 --- a/src/lightmem/configs/retriever/embeddingretriever/qdrant.py +++ b/src/lightmem/configs/retriever/embeddingretriever/qdrant.py @@ -1,6 +1,8 @@ -from pydantic import BaseModel, Field, model_validator from typing import Any, ClassVar, Dict, Optional +from pydantic import BaseModel, Field, model_validator + + class QdrantConfig(BaseModel): from qdrant_client import QdrantClient diff --git a/src/lightmem/configs/text_embedder/base.py b/src/lightmem/configs/text_embedder/base.py index b047d27..4c0c35e 100644 --- a/src/lightmem/configs/text_embedder/base.py +++ b/src/lightmem/configs/text_embedder/base.py @@ -1,7 +1,10 @@ +from typing import Any, ClassVar, Dict, List, Optional, Union + from pydantic import BaseModel, Field, model_validator -from typing import Dict, Optional, Any, List, Union, ClassVar + from lightmem.configs.text_embedder.base_config import BaseTextEmbedderConfig + class TextEmbedderConfig(BaseModel): model_name: str = Field( default="huggingface", diff --git a/src/lightmem/configs/text_embedder/base_config.py b/src/lightmem/configs/text_embedder/base_config.py index a15c243..24ea3c5 100644 --- a/src/lightmem/configs/text_embedder/base_config.py +++ b/src/lightmem/configs/text_embedder/base_config.py @@ -1,5 +1,7 @@ +from typing import Any, Dict, Optional + from pydantic import BaseModel -from typing import Optional, Dict, Any + class BaseTextEmbedderConfig(BaseModel): model: Optional[str] = None diff --git a/src/lightmem/configs/topic_segmenter/base.py b/src/lightmem/configs/topic_segmenter/base.py index 1ef0783..683f36c 100644 --- a/src/lightmem/configs/topic_segmenter/base.py +++ b/src/lightmem/configs/topic_segmenter/base.py @@ -1,5 +1,7 @@ +from typing import ClassVar, List, Optional + from pydantic import BaseModel, Field, model_validator -from typing import Dict, Optional, Type, Any, List, ClassVar + class TopicSegmenterConfig(BaseModel): model_name: str = Field(description="The TopicSegmenter model or Deployment platform (e.g., 'openai', 'deepseek', 'ollama', 'llmlingua-2')", default="llmlingua-2") diff --git a/src/lightmem/factory/memory_buffer/sensory_memory.py b/src/lightmem/factory/memory_buffer/sensory_memory.py index 8c79566..c9fcc8b 100644 --- a/src/lightmem/factory/memory_buffer/sensory_memory.py +++ b/src/lightmem/factory/memory_buffer/sensory_memory.py @@ -1,5 +1,7 @@ +from typing import Dict, List + import numpy as np -from typing import List, Dict, Optional, Any + class SenMemBufferManager: def __init__(self, max_tokens: int = 512, tokenizer = None): diff --git a/src/lightmem/factory/memory_buffer/short_term_memory.py b/src/lightmem/factory/memory_buffer/short_term_memory.py index 2bfe201..0189b1b 100644 --- a/src/lightmem/factory/memory_buffer/short_term_memory.py +++ b/src/lightmem/factory/memory_buffer/short_term_memory.py @@ -1,6 +1,8 @@ -from typing import List, Dict, Optional, Any, Union +from typing import Any, Dict, List, Optional + from lightmem.memory.utils import resolve_tokenizer + class ShortMemBufferManager: def __init__(self, max_tokens: int = 2000, tokenizer: Optional[Any] = None): self.max_tokens = max_tokens diff --git a/src/lightmem/factory/memory_manager/deepseek.py b/src/lightmem/factory/memory_manager/deepseek.py index 0515e46..6ddc019 100644 --- a/src/lightmem/factory/memory_manager/deepseek.py +++ b/src/lightmem/factory/memory_manager/deepseek.py @@ -1,8 +1,9 @@ import concurrent import json import os +from typing import Dict, List, Literal, Optional + from openai import OpenAI -from typing import List, Dict, Optional, Literal, Any from lightmem.configs.memory_manager.base_config import BaseMemoryManagerConfig from lightmem.memory.prompts import EXTRACTION_PROMPTS, METADATA_GENERATE_PROMPT diff --git a/src/lightmem/factory/memory_manager/factory.py b/src/lightmem/factory/memory_manager/factory.py index edcc890..232a9d1 100644 --- a/src/lightmem/factory/memory_manager/factory.py +++ b/src/lightmem/factory/memory_manager/factory.py @@ -1,5 +1,6 @@ -from typing import Dict, Optional from importlib import import_module +from typing import Dict + from lightmem.configs.memory_manager.base import MemoryManagerConfig diff --git a/src/lightmem/factory/memory_manager/ollama.py b/src/lightmem/factory/memory_manager/ollama.py index ed54e64..f845ac9 100644 --- a/src/lightmem/factory/memory_manager/ollama.py +++ b/src/lightmem/factory/memory_manager/ollama.py @@ -1,6 +1,6 @@ import concurrent import json -from typing import Dict, List, Optional, Literal, Any, Union +from typing import Dict, List, Literal, Optional, Union try: import ollama diff --git a/src/lightmem/factory/memory_manager/openai.py b/src/lightmem/factory/memory_manager/openai.py index 684e66f..9f77147 100644 --- a/src/lightmem/factory/memory_manager/openai.py +++ b/src/lightmem/factory/memory_manager/openai.py @@ -1,11 +1,13 @@ import concurrent -from collections import defaultdict -from openai import OpenAI -from typing import List, Dict, Optional, Literal, Any -import json, os, warnings +import json +import os +from typing import Dict, List, Literal, Optional + import httpx -from lightmem.memory.prompts import EXTRACTION_PROMPTS, METADATA_GENERATE_PROMPT +from openai import OpenAI + from lightmem.configs.memory_manager.base_config import BaseMemoryManagerConfig +from lightmem.memory.prompts import EXTRACTION_PROMPTS, METADATA_GENERATE_PROMPT from lightmem.memory.utils import clean_response model_name_context_windows = { diff --git a/src/lightmem/factory/memory_manager/transformers.py b/src/lightmem/factory/memory_manager/transformers.py index b60d93e..ca02661 100644 --- a/src/lightmem/factory/memory_manager/transformers.py +++ b/src/lightmem/factory/memory_manager/transformers.py @@ -1,9 +1,8 @@ import concurrent -import json -from typing import Dict, List, Optional, Literal, Any, Union +from typing import Dict, List, Literal, Optional import torch -from transformers import AutoTokenizer, AutoModelForCausalLM +from transformers import AutoModelForCausalLM, AutoTokenizer from lightmem.configs.memory_manager.base_config import BaseMemoryManagerConfig from lightmem.memory.prompts import EXTRACTION_PROMPTS, METADATA_GENERATE_PROMPT diff --git a/src/lightmem/factory/memory_manager/vllm.py b/src/lightmem/factory/memory_manager/vllm.py index 4f4268d..b1e624e 100644 --- a/src/lightmem/factory/memory_manager/vllm.py +++ b/src/lightmem/factory/memory_manager/vllm.py @@ -9,8 +9,9 @@ import json import os import re +from typing import Any, Dict, List, Literal, Optional + from openai import OpenAI -from typing import List, Dict, Optional, Literal, Any from lightmem.configs.memory_manager.base_config import BaseMemoryManagerConfig from lightmem.memory.prompts import EXTRACTION_PROMPTS, METADATA_GENERATE_PROMPT diff --git a/src/lightmem/factory/memory_manager/vllm_offline.py b/src/lightmem/factory/memory_manager/vllm_offline.py index bcaf998..200c731 100644 --- a/src/lightmem/factory/memory_manager/vllm_offline.py +++ b/src/lightmem/factory/memory_manager/vllm_offline.py @@ -5,11 +5,11 @@ """ import concurrent -import json import warnings -from typing import Dict, List, Optional, Literal, Any, Union +from typing import Dict, List, Literal, Optional, Union import torch + try: from vllm import LLM, SamplingParams except ImportError: diff --git a/src/lightmem/factory/pre_compressor/entropy_compress.py b/src/lightmem/factory/pre_compressor/entropy_compress.py index a227969..724f031 100644 --- a/src/lightmem/factory/pre_compressor/entropy_compress.py +++ b/src/lightmem/factory/pre_compressor/entropy_compress.py @@ -1,6 +1,8 @@ +from typing import Dict, List, Optional + import torch -from transformers import AutoTokenizer, AutoModelForCausalLM -from typing import List, Dict, Optional +from transformers import AutoModelForCausalLM, AutoTokenizer + from lightmem.configs.pre_compressor.entropy_compress import EntropyCompressorConfig diff --git a/src/lightmem/factory/pre_compressor/factory.py b/src/lightmem/factory/pre_compressor/factory.py index 8e4fef6..44f8e6d 100644 --- a/src/lightmem/factory/pre_compressor/factory.py +++ b/src/lightmem/factory/pre_compressor/factory.py @@ -1,6 +1,8 @@ -from typing import Dict, Optional from importlib import import_module -from lightmem.configs.pre_compressor.base import PreCompressorConfig +from typing import Dict + +from lightmem.configs.pre_compressor.base import PreCompressorConfig + class PreCompressorFactory: _MODEL_MAPPING: Dict[str, str] = { diff --git a/src/lightmem/factory/pre_compressor/llmlingua_2.py b/src/lightmem/factory/pre_compressor/llmlingua_2.py index daee4f5..e4169e4 100644 --- a/src/lightmem/factory/pre_compressor/llmlingua_2.py +++ b/src/lightmem/factory/pre_compressor/llmlingua_2.py @@ -1,4 +1,5 @@ -from typing import Dict, Optional, List, Union, Any +from typing import Any, Dict, List, Optional, Union + from transformers import PreTrainedTokenizerBase from lightmem.configs.pre_compressor.llmlingua_2 import LlmLingua2Config diff --git a/src/lightmem/factory/retriever/contextretriever/factory.py b/src/lightmem/factory/retriever/contextretriever/factory.py index 82f3d70..237522b 100644 --- a/src/lightmem/factory/retriever/contextretriever/factory.py +++ b/src/lightmem/factory/retriever/contextretriever/factory.py @@ -1,5 +1,6 @@ -from typing import Dict, Optional from importlib import import_module +from typing import Dict + from lightmem.configs.retriever.contextretriever.base import ContextRetrieverConfig diff --git a/src/lightmem/factory/retriever/embeddingretriever/factory.py b/src/lightmem/factory/retriever/embeddingretriever/factory.py index 3b6b793..f24f24a 100644 --- a/src/lightmem/factory/retriever/embeddingretriever/factory.py +++ b/src/lightmem/factory/retriever/embeddingretriever/factory.py @@ -1,7 +1,9 @@ -from typing import Dict, Optional from importlib import import_module +from typing import Dict + from lightmem.configs.retriever.embeddingretriever.base import EmbeddingRetrieverConfig + class EmbeddingRetrieverFactory: _MODEL_MAPPING: Dict[str, str] = { "qdrant": "lightmem.factory.retriever.embeddingretriever.qdrant.Qdrant", diff --git a/src/lightmem/factory/retriever/embeddingretriever/qdrant.py b/src/lightmem/factory/retriever/embeddingretriever/qdrant.py index b586e9b..4c28e3a 100644 --- a/src/lightmem/factory/retriever/embeddingretriever/qdrant.py +++ b/src/lightmem/factory/retriever/embeddingretriever/qdrant.py @@ -1,21 +1,22 @@ import logging import os import shutil +from typing import Any, Optional from qdrant_client import QdrantClient from qdrant_client.models import ( Distance, FieldCondition, Filter, + MatchAny, MatchValue, PointIdsList, PointStruct, Range, VectorParams, - MatchAny, ) + from lightmem.configs.retriever.embeddingretriever.qdrant import QdrantConfig -from typing import Any, Dict, Optional logger = logging.getLogger(__name__) diff --git a/src/lightmem/factory/text_embedder/factory.py b/src/lightmem/factory/text_embedder/factory.py index 57e7504..44f4f4d 100644 --- a/src/lightmem/factory/text_embedder/factory.py +++ b/src/lightmem/factory/text_embedder/factory.py @@ -1,7 +1,9 @@ -from typing import Dict, Optional from importlib import import_module +from typing import Dict + from lightmem.configs.text_embedder.base import TextEmbedderConfig + class TextEmbedderFactory: _MODEL_MAPPING: Dict[str, str] = { "huggingface": "lightmem.factory.text_embedder.huggingface.TextEmbedderHuggingface", diff --git a/src/lightmem/factory/text_embedder/huggingface.py b/src/lightmem/factory/text_embedder/huggingface.py index 67257a9..5939413 100644 --- a/src/lightmem/factory/text_embedder/huggingface.py +++ b/src/lightmem/factory/text_embedder/huggingface.py @@ -1,9 +1,12 @@ +from typing import Optional + +import numpy as np from openai import OpenAI -from typing import Optional, Literal from sentence_transformers import SentenceTransformer -import numpy as np + from lightmem.configs.text_embedder.base_config import BaseTextEmbedderConfig + class TextEmbedderHuggingface: def __init__(self, config: Optional[BaseTextEmbedderConfig] = None): self.config = config diff --git a/src/lightmem/factory/text_embedder/openai.py b/src/lightmem/factory/text_embedder/openai.py index 05be917..e5fe27f 100644 --- a/src/lightmem/factory/text_embedder/openai.py +++ b/src/lightmem/factory/text_embedder/openai.py @@ -1,7 +1,8 @@ -from openai import OpenAI -from typing import Optional, List, Union -import os +from typing import List, Optional, Union + import httpx +from openai import OpenAI + from lightmem.configs.text_embedder.base_config import BaseTextEmbedderConfig diff --git a/src/lightmem/factory/topic_segmenter/factory.py b/src/lightmem/factory/topic_segmenter/factory.py index 4c4fd4a..8c2defa 100644 --- a/src/lightmem/factory/topic_segmenter/factory.py +++ b/src/lightmem/factory/topic_segmenter/factory.py @@ -1,7 +1,9 @@ -from typing import Dict, Optional from importlib import import_module +from typing import Dict + from lightmem.configs.topic_segmenter.base import TopicSegmenterConfig + class TopicSegmenterFactory: _MODEL_MAPPING: Dict[str, str] = { "llmlingua-2": "lightmem.factory.topic_segmenter.llmlingua_2.LlmLingua2Segmenter", diff --git a/src/lightmem/factory/topic_segmenter/llmlingua_2.py b/src/lightmem/factory/topic_segmenter/llmlingua_2.py index 841c176..8e4d515 100644 --- a/src/lightmem/factory/topic_segmenter/llmlingua_2.py +++ b/src/lightmem/factory/topic_segmenter/llmlingua_2.py @@ -1,6 +1,9 @@ -from typing import Dict, Optional, List, Any -import torch, numpy as np -from transformers import AutoTokenizer, AutoModel +from typing import Any, Dict, List, Optional + +import numpy as np +import torch +from transformers import AutoModel, AutoTokenizer + class LlmLingua2Segmenter: def __init__(self, config: Optional[Dict] = None, shared: bool = False, compressor=None): diff --git a/src/lightmem/memory/graph.py b/src/lightmem/memory/graph.py index 881033e..f018c14 100644 --- a/src/lightmem/memory/graph.py +++ b/src/lightmem/memory/graph.py @@ -1,2 +1,2 @@ class GraphMem: - \ No newline at end of file + pass diff --git a/src/lightmem/memory/lightmem.py b/src/lightmem/memory/lightmem.py index f0ea388..fc71aec 100644 --- a/src/lightmem/memory/lightmem.py +++ b/src/lightmem/memory/lightmem.py @@ -1,26 +1,34 @@ -import uuid -import re -import copy import concurrent -import logging +import copy import json +import re import threading +import uuid from datetime import datetime, timedelta -from typing import Any, Dict, Literal, Optional, List, Tuple, Union +from typing import Any, Dict, List, Literal, Optional, Tuple, Union + from pydantic import ValidationError + from lightmem.configs.base import BaseMemoryConfigs -from lightmem.factory.pre_compressor.factory import PreCompressorFactory -from lightmem.factory.topic_segmenter.factory import TopicSegmenterFactory -from lightmem.factory.memory_manager.factory import MemoryManagerFactory -from lightmem.factory.text_embedder.factory import TextEmbedderFactory -from lightmem.factory.retriever.contextretriever.factory import ContextRetrieverFactory -from lightmem.factory.retriever.embeddingretriever.factory import EmbeddingRetrieverFactory -from lightmem.factory.retriever.embeddingretriever.qdrant import QdrantConfig +from lightmem.configs.logging.utils import get_logger from lightmem.factory.memory_buffer.sensory_memory import SenMemBufferManager from lightmem.factory.memory_buffer.short_term_memory import ShortMemBufferManager -from lightmem.memory.utils import MemoryEntry, assign_sequence_numbers_with_timestamps, save_memory_entries,convert_extraction_results_to_memory_entries,normalize_extraction_prompts,process_extraction_results -from lightmem.memory.prompts import METADATA_GENERATE_PROMPT, UPDATE_PROMPT -from lightmem.configs.logging.utils import get_logger +from lightmem.factory.memory_manager.factory import MemoryManagerFactory +from lightmem.factory.pre_compressor.factory import PreCompressorFactory +from lightmem.factory.retriever.contextretriever.factory import ContextRetrieverFactory +from lightmem.factory.retriever.embeddingretriever.factory import ( + EmbeddingRetrieverFactory, +) +from lightmem.factory.text_embedder.factory import TextEmbedderFactory +from lightmem.factory.topic_segmenter.factory import TopicSegmenterFactory +from lightmem.memory.prompts import UPDATE_PROMPT +from lightmem.memory.utils import ( + assign_sequence_numbers_with_timestamps, + convert_extraction_results_to_memory_entries, + normalize_extraction_prompts, + process_extraction_results, + save_memory_entries, +) GLOBAL_TOPIC_IDX = 0 GLOBAL_LAST_SUMMARY_TIME = None @@ -723,18 +731,18 @@ def summarize( top_k_seeds: int = 15, ) -> Dict: from lightmem.memory.utils import ( - initialize_time_pointer, + build_batch_result, + build_empty_result, + build_single_result, + build_summary_item, + call_summary_llm, + check_has_more_entries, + format_entries_for_prompt, get_window_entries, + initialize_time_pointer, mark_entries_and_get_next_time, - check_has_more_entries, retrieve_supplementary_entries, - format_entries_for_prompt, - call_summary_llm, store_summary, - build_summary_item, - build_single_result, - build_batch_result, - build_empty_result ) global GLOBAL_LAST_SUMMARY_TIME diff --git a/src/lightmem/memory/utils.py b/src/lightmem/memory/utils.py index b29539b..e851462 100644 --- a/src/lightmem/memory/utils.py +++ b/src/lightmem/memory/utils.py @@ -1,14 +1,15 @@ +import json import os import re -import json -from datetime import datetime -from typing import List, Dict, Literal, Optional, Any, Tuple, Union -import tiktoken import uuid from dataclasses import dataclass, field -from transformers.tokenization_utils_fast import PreTrainedTokenizerFast +from datetime import datetime +from typing import Any, Dict, List, Literal, Optional, Tuple, Union + +import tiktoken from transformers.tokenization_utils import PreTrainedTokenizer -from typing import Optional, Union, Dict +from transformers.tokenization_utils_fast import PreTrainedTokenizerFast + @dataclass class MemoryEntry: @@ -57,9 +58,9 @@ def clean_response(response: str) -> List[Dict[str, Any]]: def assign_sequence_numbers_with_timestamps(extract_list, offset_ms: int = 500, topic_id_mapping: List[List[int]] = None): - from datetime import datetime, timedelta - from collections import defaultdict import re + from collections import defaultdict + from datetime import datetime, timedelta current_index = 0 timestamps_list = [] diff --git a/src/lightmem/memory_toolkits/inference_utils/backends.py b/src/lightmem/memory_toolkits/inference_utils/backends.py index 46fe652..e583aba 100644 --- a/src/lightmem/memory_toolkits/inference_utils/backends.py +++ b/src/lightmem/memory_toolkits/inference_utils/backends.py @@ -1,14 +1,16 @@ -from openai import OpenAI -import warnings +import warnings from concurrent.futures import ThreadPoolExecutor from typing import ( - Optional, - Dict, - Any, - Callable, - List, + Any, + Callable, + Dict, + List, + Optional, ) +from openai import OpenAI + + class OpenAIClient(OpenAI): """ A subclass of ``openai.OpenAI`` that wraps around the chat completions API to provide a simple interface @@ -157,7 +159,7 @@ def __call__( ) if len(kwargs) > 0: - from vllm import SamplingParams + from vllm import SamplingParams sampling_params = SamplingParams(**kwargs) outputs = self.model.generate(texts, sampling_params) else: diff --git a/src/lightmem/memory_toolkits/inference_utils/base_operator.py b/src/lightmem/memory_toolkits/inference_utils/base_operator.py index 99aa64d..94b2f26 100644 --- a/src/lightmem/memory_toolkits/inference_utils/base_operator.py +++ b/src/lightmem/memory_toolkits/inference_utils/base_operator.py @@ -1,19 +1,23 @@ from __future__ import annotations -from .prompts import get_prompt -from string import Template + import inspect +import os +import time from abc import ABC, abstractmethod -from .backends import get_interface_for_inference -from tqdm import tqdm -import time -import os +from string import Template from typing import ( - Any, - List, - Optional, - Dict, + Any, + Dict, + List, + Optional, ) +from tqdm import tqdm + +from .backends import get_interface_for_inference +from .prompts import get_prompt + + class NonCachedLLMOperator(ABC): def __init__( diff --git a/src/lightmem/memory_toolkits/inference_utils/operators.py b/src/lightmem/memory_toolkits/inference_utils/operators.py index d24c2ea..c0d84fc 100644 --- a/src/lightmem/memory_toolkits/inference_utils/operators.py +++ b/src/lightmem/memory_toolkits/inference_utils/operators.py @@ -1,14 +1,17 @@ -from .base_operator import NonCachedLLMOperator -import numpy as np import json import re -from typing import ( - List, - Dict, - Any, - Optional, +from typing import ( + Any, + Dict, + List, + Optional, ) +import numpy as np + +from .base_operator import NonCachedLLMOperator + + class QuestionAnsweringOperator(NonCachedLLMOperator): """An operator for question answering.""" diff --git a/src/lightmem/memory_toolkits/memories/__init__.py b/src/lightmem/memory_toolkits/memories/__init__.py index 9e7d970..d5695fd 100644 --- a/src/lightmem/memory_toolkits/memories/__init__.py +++ b/src/lightmem/memory_toolkits/memories/__init__.py @@ -2,8 +2,8 @@ from collections import OrderedDict from collections.abc import Iterator from typing import ( - Any, - Type, + Any, + Type, Union, ) diff --git a/src/lightmem/memory_toolkits/memories/datasets/base.py b/src/lightmem/memory_toolkits/memories/datasets/base.py index 363df72..f4ae943 100644 --- a/src/lightmem/memory_toolkits/memories/datasets/base.py +++ b/src/lightmem/memory_toolkits/memories/datasets/base.py @@ -1,25 +1,27 @@ from __future__ import annotations -from pydantic import ( - BaseModel, - Field, - model_validator, - field_serializer, - field_validator, - ConfigDict, -) + +import random from abc import ABC, abstractmethod -from types import MappingProxyType from datetime import datetime from functools import total_ordering -import random +from types import MappingProxyType from typing import ( - List, - Any, + Any, + Dict, Iterator, + List, + Mapping, + Optional, Tuple, - Mapping, - Dict, - Optional, +) + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_serializer, + field_validator, + model_validator, ) TIMESTAMP_FORMAT = "%Y-%m-%d (%a) %H:%M" diff --git a/src/lightmem/memory_toolkits/memories/datasets/locomo.py b/src/lightmem/memory_toolkits/memories/datasets/locomo.py index 64c8509..1f924de 100644 --- a/src/lightmem/memory_toolkits/memories/datasets/locomo.py +++ b/src/lightmem/memory_toolkits/memories/datasets/locomo.py @@ -1,22 +1,22 @@ from __future__ import annotations -import os import base64 -import requests -from io import BytesIO -from PIL import Image -import openai - import json +import os from datetime import datetime -from typing import Dict, Any, List, Optional, Tuple +from io import BytesIO +from typing import Any, Dict, List, Optional, Tuple + +import openai +import requests +from PIL import Image from .base import ( MemoryDataset, - Trajectory, - Session, - QuestionAnswerPair, Message, + QuestionAnswerPair, + Session, + Trajectory, ) diff --git a/src/lightmem/memory_toolkits/memories/datasets/longmemeval.py b/src/lightmem/memory_toolkits/memories/datasets/longmemeval.py index fb69612..d93432e 100644 --- a/src/lightmem/memory_toolkits/memories/datasets/longmemeval.py +++ b/src/lightmem/memory_toolkits/memories/datasets/longmemeval.py @@ -1,14 +1,16 @@ from __future__ import annotations -import json + +import json +from datetime import datetime +from typing import Any, Dict, Tuple + from .base import ( - MemoryDataset, - Trajectory, - Session, - QuestionAnswerPair, - Message, + MemoryDataset, + Message, + QuestionAnswerPair, + Session, + Trajectory, ) -from datetime import datetime -from typing import Dict, Any, Tuple class LongMemEval(MemoryDataset): diff --git a/src/lightmem/memory_toolkits/memories/layers/amem.py b/src/lightmem/memory_toolkits/memories/layers/amem.py index fa1aa61..4286682 100644 --- a/src/lightmem/memory_toolkits/memories/layers/amem.py +++ b/src/lightmem/memory_toolkits/memories/layers/amem.py @@ -1,25 +1,29 @@ from __future__ import annotations -from .base import BaseMemoryLayer + +import json +import os +import pickle +from typing import ( + Any, + Dict, + List, + Literal, + Optional, +) + from pydantic import ( - BaseModel, - Field, + BaseModel, + Field, model_validator, ) + +from .base import BaseMemoryLayer from .baselines.agentic_memory.memory_system import ( - AgenticMemorySystem, - MemoryNote, -) -import pickle -import os -import json -from typing import ( - Literal, - List, - Dict, - Any, - Optional, + AgenticMemorySystem, + MemoryNote, ) + class AMEMConfig(BaseModel): """The default configuration for A-MEM.""" diff --git a/src/lightmem/memory_toolkits/memories/layers/base.py b/src/lightmem/memory_toolkits/memories/layers/base.py index c38289f..754ddaa 100644 --- a/src/lightmem/memory_toolkits/memories/layers/base.py +++ b/src/lightmem/memory_toolkits/memories/layers/base.py @@ -1,10 +1,10 @@ from abc import ABC, abstractmethod from typing import ( - Any, - Dict, - List, - Union, + Any, + Dict, + List, Optional, + Union, ) # Note: This base interface will evolve as additional memory systems are integrated. diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/llm_controller.py b/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/llm_controller.py index c9eb8b2..e8654a5 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/llm_controller.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/llm_controller.py @@ -1,9 +1,11 @@ -from typing import Dict, Optional, Literal, Any -import os import json +import os from abc import ABC, abstractmethod +from typing import Any, Literal, Optional + from litellm import completion + class BaseLLMController(ABC): @abstractmethod def get_completion(self, prompt: str) -> str: @@ -42,7 +44,6 @@ def get_completion(self, prompt: str, response_format: dict, temperature: float class OllamaController(BaseLLMController): def __init__(self, model: str = "llama2"): - from ollama import chat self.model = model def _generate_empty_value(self, schema_type: str, schema_items: dict = None) -> Any: diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/memory_system.py b/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/memory_system.py index 4622f91..a3657c5 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/memory_system.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/memory_system.py @@ -1,23 +1,11 @@ -import keyword -from typing import List, Dict, Optional, Any, Tuple +import json +import logging import uuid from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + from .llm_controller import LLMController from .retrievers import ChromaRetriever -import json -import logging -from rank_bm25 import BM25Okapi -from sentence_transformers import SentenceTransformer -import numpy as np -from sklearn.metrics.pairwise import cosine_similarity -import os -from abc import ABC, abstractmethod -from transformers import AutoModel, AutoTokenizer -from nltk.tokenize import word_tokenize -import pickle -from pathlib import Path -from litellm import completion -import time logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/retrievers.py b/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/retrievers.py index d91e590..3367b32 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/retrievers.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/agentic_memory/retrievers.py @@ -1,19 +1,15 @@ -from typing import List, Dict, Any, Optional, Union -from sentence_transformers import SentenceTransformer -from rank_bm25 import BM25Okapi -import nltk -import numpy as np -from sklearn.metrics.pairwise import cosine_similarity +import json +import os +from typing import Dict, Optional + import chromadb from chromadb.config import Settings -import pickle -from nltk.tokenize import word_tokenize -import os -import json from chromadb.utils.embedding_functions import ( + OpenAIEmbeddingFunction, SentenceTransformerEmbeddingFunction, - OpenAIEmbeddingFunction ) +from nltk.tokenize import word_tokenize + def simple_tokenize(text): return word_tokenize(text) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/main.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/main.py index 0e257a0..999a97b 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/main.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/main.py @@ -6,9 +6,9 @@ import httpx import requests - from mem0.client.project import AsyncProject, Project from mem0.client.utils import api_error_handler + # Exception classes are referenced in docstrings only from mem0.memory.setup import get_user_id, setup_config from mem0.memory.telemetry import capture_client_event diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/project.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/project.py index 45c9081..00a38b9 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/project.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/project.py @@ -3,10 +3,10 @@ from typing import Any, Dict, List, Optional import httpx -from pydantic import BaseModel, ConfigDict, Field - from mem0.client.utils import api_error_handler from mem0.memory.telemetry import capture_client_event +from pydantic import BaseModel, ConfigDict, Field + # Exception classes are referenced in docstrings only logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/utils.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/utils.py index 06a1c0e..fcebffa 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/utils.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/client/utils.py @@ -1,7 +1,7 @@ import json import logging -import httpx +import httpx from mem0.exceptions import ( NetworkError, create_exception_from_response, diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/base.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/base.py index 147d259..27ba6c7 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/base.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/base.py @@ -1,12 +1,11 @@ import os from typing import Any, Dict, Optional -from pydantic import BaseModel, Field - from mem0.embeddings.configs import EmbedderConfig from mem0.graphs.configs import GraphStoreConfig from mem0.llms.configs import LlmConfig from mem0.vector_stores.configs import VectorStoreConfig +from pydantic import BaseModel, Field # Set up the directory path home_dir = os.path.expanduser("~") diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/embeddings/base.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/embeddings/base.py index 0737088..088026d 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/embeddings/base.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/embeddings/base.py @@ -3,7 +3,6 @@ from typing import Dict, Optional, Union import httpx - from mem0.configs.base import AzureConfig diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/vector_stores/databricks.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/vector_stores/databricks.py index 6af0664..be4b093 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/vector_stores/databricks.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/configs/vector_stores/databricks.py @@ -1,9 +1,12 @@ from typing import Any, Dict, Optional +from databricks.sdk.service.vectorsearch import ( + EndpointType, + PipelineType, + VectorIndexType, +) from pydantic import BaseModel, ConfigDict, Field, model_validator -from databricks.sdk.service.vectorsearch import EndpointType, VectorIndexType, PipelineType - class DatabricksConfig(BaseModel): """Configuration for Databricks Vector Search vector store.""" diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/aws_bedrock.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/aws_bedrock.py index 5c3c1ac..a0c608a 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/aws_bedrock.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/aws_bedrock.py @@ -8,7 +8,6 @@ raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.") import numpy as np - from mem0.configs.embeddings.base import BaseEmbedderConfig from mem0.embeddings.base import EmbeddingBase diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/azure_openai.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/azure_openai.py index 547ec0c..faf8363 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/azure_openai.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/azure_openai.py @@ -2,10 +2,9 @@ from typing import Literal, Optional from azure.identity import DefaultAzureCredential, get_bearer_token_provider -from openai import AzureOpenAI - from mem0.configs.embeddings.base import BaseEmbedderConfig from mem0.embeddings.base import EmbeddingBase +from openai import AzureOpenAI SCOPE = "https://cognitiveservices.azure.com/.default" diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/gemini.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/gemini.py index 203b311..c481bf1 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/gemini.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/gemini.py @@ -3,7 +3,6 @@ from google import genai from google.genai import types - from mem0.configs.embeddings.base import BaseEmbedderConfig from mem0.embeddings.base import EmbeddingBase diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/huggingface.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/huggingface.py index 934c69a..8caab88 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/huggingface.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/embeddings/huggingface.py @@ -1,11 +1,10 @@ import logging from typing import Literal, Optional -from openai import OpenAI -from sentence_transformers import SentenceTransformer - from mem0.configs.embeddings.base import BaseEmbedderConfig from mem0.embeddings.base import EmbeddingBase +from openai import OpenAI +from sentence_transformers import SentenceTransformer logging.getLogger("transformers").setLevel(logging.WARNING) logging.getLogger("sentence_transformers").setLevel(logging.WARNING) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunedb.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunedb.py index d0ca68f..9c545ab 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunedb.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunedb.py @@ -1,6 +1,7 @@ import logging import uuid from datetime import datetime + import pytz from .base import NeptuneBase diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunegraph.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunegraph.py index c926448..a6d4e22 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunegraph.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/graphs/neptune/neptunegraph.py @@ -3,8 +3,8 @@ from .base import NeptuneBase try: - from langchain_aws import NeptuneAnalyticsGraph from botocore.config import Config + from langchain_aws import NeptuneAnalyticsGraph except ImportError: raise ImportError("langchain_aws is not installed. Please install it using 'make install_all'.") diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/aws_bedrock.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/aws_bedrock.py index b29cf59..2128dce 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/aws_bedrock.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/aws_bedrock.py @@ -9,8 +9,8 @@ except ImportError: raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.") -from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.llms.aws_bedrock import AWSBedrockConfig +from mem0.configs.llms.base import BaseLlmConfig from mem0.llms.base import LLMBase logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai.py index 6ddb50b..803787b 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai.py @@ -3,12 +3,11 @@ from typing import Dict, List, Optional, Union from azure.identity import DefaultAzureCredential, get_bearer_token_provider -from openai import AzureOpenAI - from mem0.configs.llms.azure import AzureOpenAIConfig from mem0.configs.llms.base import BaseLlmConfig from mem0.llms.base import LLMBase from mem0.memory.utils import extract_json +from openai import AzureOpenAI SCOPE = "https://cognitiveservices.azure.com/.default" diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai_structured.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai_structured.py index fd2bae0..0d2cc4c 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai_structured.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/azure_openai_structured.py @@ -2,10 +2,9 @@ from typing import Dict, List, Optional from azure.identity import DefaultAzureCredential, get_bearer_token_provider -from openai import AzureOpenAI - from mem0.configs.llms.base import BaseLlmConfig from mem0.llms.base import LLMBase +from openai import AzureOpenAI SCOPE = "https://cognitiveservices.azure.com/.default" diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/deepseek.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/deepseek.py index a987706..39f280f 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/deepseek.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/deepseek.py @@ -2,12 +2,11 @@ import os from typing import Dict, List, Optional, Union -from openai import OpenAI - from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.llms.deepseek import DeepSeekConfig from mem0.llms.base import LLMBase from mem0.memory.utils import extract_json +from openai import OpenAI class DeepSeekLLM(LLMBase): diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/lmstudio.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/lmstudio.py index aab5d07..30a7e47 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/lmstudio.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/lmstudio.py @@ -1,12 +1,11 @@ import json from typing import Dict, List, Optional, Union -from openai import OpenAI - from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.llms.lmstudio import LMStudioConfig from mem0.llms.base import LLMBase from mem0.memory.utils import extract_json +from openai import OpenAI class LMStudioLLM(LLMBase): diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/openai.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/openai.py index b6ad538..eb86ad5 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/openai.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/openai.py @@ -3,12 +3,11 @@ import os from typing import Dict, List, Optional, Union -from openai import OpenAI - from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.llms.openai import OpenAIConfig from mem0.llms.base import LLMBase from mem0.memory.utils import extract_json +from openai import OpenAI class OpenAILLM(LLMBase): diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/sarvam.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/sarvam.py index 6ef836e..8fdc034 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/sarvam.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/sarvam.py @@ -2,7 +2,6 @@ from typing import Dict, List, Optional import requests - from mem0.configs.llms.base import BaseLlmConfig from mem0.llms.base import LLMBase diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/vllm.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/vllm.py index f7cbfbc..108676a 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/vllm.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/llms/vllm.py @@ -2,12 +2,11 @@ import os from typing import Dict, List, Optional, Union -from openai import OpenAI - from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.llms.vllm import VllmConfig from mem0.llms.base import LLMBase from mem0.memory.utils import extract_json +from openai import OpenAI class VllmLLM(LLMBase): diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/main.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/main.py index 7971956..f0a4217 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/main.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/main.py @@ -12,8 +12,6 @@ from typing import Any, Dict, Optional import pytz -from pydantic import ValidationError - from mem0.configs.base import MemoryConfig, MemoryItem from mem0.configs.enums import MemoryType from mem0.configs.prompts import ( @@ -38,6 +36,7 @@ LlmFactory, VectorStoreFactory, ) +from pydantic import ValidationError # Suppress SWIG deprecation warnings globally warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*SwigPy.*") diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/telemetry.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/telemetry.py index 6d822ce..0c5d6c2 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/telemetry.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/memory/telemetry.py @@ -3,10 +3,9 @@ import platform import sys -from posthog import Posthog - import mem0 from mem0.memory.setup import get_or_create_user_id +from posthog import Posthog MEM0_TELEMETRY = os.environ.get("MEM0_TELEMETRY", "True") PROJECT_API_KEY = "phc_hgJkUVJFYtmaJqrvf6CYN67TIQ8yhXAkWzUn9AMU4yX" diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/proxy/main.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/proxy/main.py index 4baaf5e..779743d 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/proxy/main.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/proxy/main.py @@ -5,7 +5,6 @@ from typing import List, Optional, Union import httpx - import mem0 try: diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_ai_search.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_ai_search.py index 6165efc..2e7956e 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_ai_search.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_ai_search.py @@ -3,10 +3,9 @@ import re from typing import List, Optional -from pydantic import BaseModel - from mem0.memory.utils import extract_json from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel try: from azure.core.credentials import AzureKeyCredential diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_mysql.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_mysql.py index 2d9ab37..9386c1c 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_mysql.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/azure_mysql.py @@ -7,8 +7,8 @@ try: import pymysql - from pymysql.cursors import DictCursor from dbutils.pooled_db import PooledDB + from pymysql.cursors import DictCursor except ImportError: raise ImportError( "Azure MySQL vector store requires PyMySQL and DBUtils. " diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/baidu.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/baidu.py index 2c211ab..95cd054 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/baidu.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/baidu.py @@ -2,9 +2,8 @@ import time from typing import Dict, Optional -from pydantic import BaseModel - from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel try: import pymochow diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/databricks.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/databricks.py index 6b5660e..3fb9006 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/databricks.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/databricks.py @@ -1,21 +1,28 @@ import json import logging import uuid -from typing import Optional, List -from datetime import datetime, date -from databricks.sdk.service.catalog import ColumnInfo, ColumnTypeName, TableType, DataSourceFormat -from databricks.sdk.service.catalog import TableConstraint, PrimaryKeyConstraint +from datetime import date, datetime +from typing import List, Optional + from databricks.sdk import WorkspaceClient +from databricks.sdk.service.catalog import ( + ColumnInfo, + ColumnTypeName, + DataSourceFormat, + PrimaryKeyConstraint, + TableConstraint, + TableType, +) from databricks.sdk.service.vectorsearch import ( - VectorIndexType, DeltaSyncVectorIndexSpecRequest, DirectAccessVectorIndexSpec, EmbeddingSourceColumn, EmbeddingVectorColumn, + VectorIndexType, ) -from pydantic import BaseModel from mem0.memory.utils import extract_json from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/elasticsearch.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/elasticsearch.py index b73eedc..8be29c0 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/elasticsearch.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/elasticsearch.py @@ -7,10 +7,9 @@ except ImportError: raise ImportError("Elasticsearch requires extra dependencies. Install with `pip install elasticsearch`") from None -from pydantic import BaseModel - from mem0.configs.vector_stores.elasticsearch import ElasticsearchConfig from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/faiss.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/faiss.py index 141df5e..28ea5db 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/faiss.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/faiss.py @@ -2,14 +2,13 @@ import os import pickle import uuid +import warnings from pathlib import Path from typing import Dict, List, Optional import numpy as np from pydantic import BaseModel -import warnings - try: # Suppress SWIG deprecation warnings from FAISS warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*SwigPy.*") diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/milvus.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/milvus.py index 41c1a33..449f534 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/milvus.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/milvus.py @@ -1,10 +1,9 @@ import logging from typing import Dict, Optional -from pydantic import BaseModel - from mem0.configs.vector_stores.milvus import MetricType from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel try: import pymilvus # noqa: F401 diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/opensearch.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/opensearch.py index 7d41757..a2fc3ea 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/opensearch.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/opensearch.py @@ -7,10 +7,9 @@ except ImportError: raise ImportError("OpenSearch requires extra dependencies. Install with `pip install opensearch-py`") from None -from pydantic import BaseModel - from mem0.configs.vector_stores.opensearch import OpenSearchConfig from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/qdrant.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/qdrant.py index 59ee9a9..ede7038 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/qdrant.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/qdrant.py @@ -2,6 +2,7 @@ import os import shutil +from mem0.vector_stores.base import VectorStoreBase from qdrant_client import QdrantClient from qdrant_client.models import ( Distance, @@ -14,8 +15,6 @@ VectorParams, ) -from mem0.vector_stores.base import VectorStoreBase - logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/redis.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/redis.py index 7fb1ada..1ac4256 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/redis.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/redis.py @@ -6,14 +6,13 @@ import numpy as np import pytz import redis +from mem0.memory.utils import extract_json +from mem0.vector_stores.base import VectorStoreBase from redis.commands.search.query import Query from redisvl.index import SearchIndex from redisvl.query import VectorQuery from redisvl.query.filter import Tag -from mem0.memory.utils import extract_json -from mem0.vector_stores.base import VectorStoreBase - logger = logging.getLogger(__name__) # TODO: Improve as these are not the best fields for the Redis's perspective. Might do away with them. diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/s3_vectors.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/s3_vectors.py index f6504c3..f4db22a 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/s3_vectors.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/s3_vectors.py @@ -2,9 +2,8 @@ import logging from typing import Dict, List, Optional -from pydantic import BaseModel - from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel try: import boto3 diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/upstash_vector.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/upstash_vector.py index 82dc0f4..7b5b5c7 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/upstash_vector.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/upstash_vector.py @@ -1,9 +1,8 @@ import logging from typing import Dict, List, Optional -from pydantic import BaseModel - from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel try: from upstash_vector import Index diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/valkey.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/valkey.py index c4539dc..5ffa3f7 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/valkey.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/valkey.py @@ -6,11 +6,10 @@ import numpy as np import pytz import valkey -from pydantic import BaseModel -from valkey.exceptions import ResponseError - from mem0.memory.utils import extract_json from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel +from valkey.exceptions import ResponseError logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/vertex_ai_vector_search.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/vertex_ai_vector_search.py index 39aa992..bd0c830 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/vertex_ai_vector_search.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/vertex_ai_vector_search.py @@ -10,12 +10,11 @@ ) from google.oauth2 import service_account from langchain.schema import Document -from pydantic import BaseModel - from mem0.configs.vector_stores.vertex_ai_vector_search import ( GoogleMatchingEngineConfig, ) from mem0.vector_stores.base import VectorStoreBase +from pydantic import BaseModel # Configure logging logging.basicConfig(level=logging.DEBUG) diff --git a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/weaviate.py b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/weaviate.py index cb1ed6d..cc77463 100644 --- a/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/weaviate.py +++ b/src/lightmem/memory_toolkits/memories/layers/baselines/mem0/vector_stores/weaviate.py @@ -13,12 +13,11 @@ ) import weaviate.classes.config as wvcc +from mem0.vector_stores.base import VectorStoreBase from weaviate.classes.init import AdditionalConfig, Auth, Timeout from weaviate.classes.query import Filter, MetadataQuery from weaviate.util import get_valid_uuid -from mem0.vector_stores.base import VectorStoreBase - logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/full_context.py b/src/lightmem/memory_toolkits/memories/layers/full_context.py index bbf5da9..4a4a08f 100644 --- a/src/lightmem/memory_toolkits/memories/layers/full_context.py +++ b/src/lightmem/memory_toolkits/memories/layers/full_context.py @@ -1,15 +1,17 @@ from __future__ import annotations -import os + import json +import os import pickle import time -from typing import Any, Dict, List, Optional, Union, Literal +from typing import Any, Dict, List, Literal, Optional, Union +from litellm import token_counter as litellm_token_counter from pydantic import BaseModel, Field, model_validator +from token_monitor import get_tokenizer_for_model from .base import BaseMemoryLayer -from token_monitor import get_tokenizer_for_model -from litellm import token_counter as litellm_token_counter + class FullContextConfig(BaseModel): """Configuration: Full-Text Memory layer (simple full-context storage and retrieval). diff --git a/src/lightmem/memory_toolkits/memories/layers/langmem.py b/src/lightmem/memory_toolkits/memories/layers/langmem.py index 6d3ea07..0385e1d 100644 --- a/src/lightmem/memory_toolkits/memories/layers/langmem.py +++ b/src/lightmem/memory_toolkits/memories/layers/langmem.py @@ -1,26 +1,29 @@ from __future__ import annotations -from .base import BaseMemoryLayer -from langgraph.store.memory import InMemoryStore + +import json +import os +import pickle +from copy import deepcopy +from typing import ( + Any, + Dict, + List, + Optional, +) + from langchain.chat_models import init_chat_model from langchain_core.language_models import BaseChatModel -from .baselines.langmem import create_memory_store_manager - +from langgraph.store.memory import InMemoryStore from pydantic import ( - BaseModel, - Field, - model_validator, + BaseModel, + Field, field_validator, + model_validator, ) -from copy import deepcopy -import pickle -import os -import json -from typing import ( - List, - Dict, - Any, - Optional, -) + +from .base import BaseMemoryLayer +from .baselines.langmem import create_memory_store_manager + class LangMemConfig(BaseModel): """The default configuration for LangMem.""" diff --git a/src/lightmem/memory_toolkits/memories/layers/memzero.py b/src/lightmem/memory_toolkits/memories/layers/memzero.py index 52bec5a..f4639a6 100644 --- a/src/lightmem/memory_toolkits/memories/layers/memzero.py +++ b/src/lightmem/memory_toolkits/memories/layers/memzero.py @@ -1,28 +1,19 @@ from __future__ import annotations -from .base import BaseMemoryLayer -from pydantic import ( - BaseModel, - Field, - model_validator, -) -from .baselines.mem0 import ( - Memory, - MemoryConfig -) -from typing import ( - Literal, - List, - Dict, - Any, - Optional, - Union -) -import os + import json -import pickle import logging +import os +import pickle +from typing import Any, Dict, List, Literal, Optional, Union +from pydantic import ( + BaseModel, + Field, + model_validator, +) +from .base import BaseMemoryLayer +from .baselines.mem0 import Memory logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memories/layers/naive_rag.py b/src/lightmem/memory_toolkits/memories/layers/naive_rag.py index ce9a7ab..e43cf76 100644 --- a/src/lightmem/memory_toolkits/memories/layers/naive_rag.py +++ b/src/lightmem/memory_toolkits/memories/layers/naive_rag.py @@ -1,20 +1,17 @@ from __future__ import annotations -import os import json import logging -from typing import Any, Dict, List, Optional, Literal, Union +import os +from collections.abc import Mapping +from types import MappingProxyType +from typing import Any, Dict, List, Literal, Optional, Union +from mem0.memory.main import Memory # type: ignore from pydantic import BaseModel, Field, model_validator from .base import BaseMemoryLayer -from mem0.memory.main import Memory # type: ignore - -from collections.abc import Mapping -from types import MappingProxyType - - logger = logging.getLogger(__name__) diff --git a/src/lightmem/memory_toolkits/memory_construction.py b/src/lightmem/memory_toolkits/memory_construction.py index 8b0fc35..6cbd4a1 100644 --- a/src/lightmem/memory_toolkits/memory_construction.py +++ b/src/lightmem/memory_toolkits/memory_construction.py @@ -1,43 +1,41 @@ import argparse import json +import os +import threading +import time from concurrent.futures import ThreadPoolExecutor, as_completed +from contextlib import nullcontext +from copy import deepcopy from datetime import datetime -from tqdm import tqdm -import time -import threading +from importlib import import_module from importlib.util import find_spec -from copy import deepcopy -import os -from contextlib import nullcontext - -from token_monitor import ( - CostStateManager, - token_monitor, - CostState, - get_tokenizer_for_model, -) -from monkey_patch import ( - PatchSpec, - MonkeyPatcher, - make_attr_patch, +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Tuple, ) from memories import ( CONFIG_MAPPING, - MEMORY_LAYERS_MAPPING, DATASET_MAPPING, + MEMORY_LAYERS_MAPPING, ) -from memories.datasets.base import Trajectory -from typing import ( - Dict, - Any, - Optional, - Tuple, - List, - Callable, +from memories.datasets.base import Message, QuestionAnswerPair, Session, Trajectory +from monkey_patch import ( + MonkeyPatcher, + PatchSpec, + make_attr_patch, ) -from memories.datasets.base import Trajectory, Message, QuestionAnswerPair, Session -from importlib import import_module +from token_monitor import ( + CostState, + CostStateManager, + get_tokenizer_for_model, + token_monitor, +) +from tqdm import tqdm _LOCK = threading.Lock() @@ -50,10 +48,10 @@ def _normalize_langmem_messages(*args, **kwargs) -> Dict[str, List[Dict[str, str """A helper function to process the messages of LangMem.""" _check_langchain_core_imports() from langchain_core.messages import ( - HumanMessage, - SystemMessage, - AIMessage, - ToolMessage, + AIMessage, + HumanMessage, + SystemMessage, + ToolMessage, ) messages = kwargs.get("messages", args[0]) diff --git a/src/lightmem/memory_toolkits/memory_evaluation.py b/src/lightmem/memory_toolkits/memory_evaluation.py index 144edb9..8415cea 100644 --- a/src/lightmem/memory_toolkits/memory_evaluation.py +++ b/src/lightmem/memory_toolkits/memory_evaluation.py @@ -1,21 +1,22 @@ -from memories.datasets.base import QuestionAnswerPair, MemoryDataset -from inference_utils.operators import ( - QuestionAnsweringOperator, - LLMExactMatch, -) -from memories import DATASET_MAPPING -import numpy as np import argparse -import json -import os +import json +import os from typing import ( - List, - Dict, Any, + Dict, + List, Optional, Tuple, ) +import numpy as np +from inference_utils.operators import ( + LLMExactMatch, + QuestionAnsweringOperator, +) +from memories import DATASET_MAPPING +from memories.datasets.base import QuestionAnswerPair + def _build_context_text(retrieved_memories: List[Dict[str, Any]]) -> str: contents = [] diff --git a/src/lightmem/memory_toolkits/memory_search.py b/src/lightmem/memory_toolkits/memory_search.py index bf182a0..51e5cf7 100644 --- a/src/lightmem/memory_toolkits/memory_search.py +++ b/src/lightmem/memory_toolkits/memory_search.py @@ -1,22 +1,23 @@ import argparse import json -from concurrent.futures import ThreadPoolExecutor, as_completed -from tqdm import tqdm import threading +from concurrent.futures import ThreadPoolExecutor, as_completed from copy import deepcopy -from memories import ( - CONFIG_MAPPING, - MEMORY_LAYERS_MAPPING, - DATASET_MAPPING, -) -from memories.datasets.base import QuestionAnswerPair, MemoryDataset from typing import ( - Dict, Any, - Optional, + Dict, List, + Optional, ) +from memories import ( + CONFIG_MAPPING, + DATASET_MAPPING, + MEMORY_LAYERS_MAPPING, +) +from memories.datasets.base import MemoryDataset, QuestionAnswerPair +from tqdm import tqdm + _LOCK = threading.Lock() diff --git a/src/lightmem/memory_toolkits/monkey_patch.py b/src/lightmem/memory_toolkits/monkey_patch.py index 4e08006..185836b 100644 --- a/src/lightmem/memory_toolkits/monkey_patch.py +++ b/src/lightmem/memory_toolkits/monkey_patch.py @@ -60,16 +60,18 @@ """ from __future__ import annotations + from dataclasses import dataclass -from pydantic import BaseModel from typing import ( - Callable, - Any, - List, - Dict, - Tuple, + Any, + Callable, + Dict, + List, + Tuple, ) +from pydantic import BaseModel + @dataclass class PatchSpec: diff --git a/src/lightmem/memory_toolkits/token_monitor.py b/src/lightmem/memory_toolkits/token_monitor.py index a466aea..9e761e1 100644 --- a/src/lightmem/memory_toolkits/token_monitor.py +++ b/src/lightmem/memory_toolkits/token_monitor.py @@ -1,24 +1,27 @@ from __future__ import annotations + import functools +import inspect import json from collections import deque -from threading import RLock, Lock -from litellm import token_counter as litellm_token_counter -from litellm.types.utils import SelectTokenizerResponse -from litellm import encoding -import tiktoken -from tokenizers import Tokenizer from datetime import datetime -import inspect +from threading import Lock, RLock from typing import ( - Any, - Callable, - Dict, - List, - Optional, + Any, + Callable, + Dict, + List, + Optional, Tuple, ) +import tiktoken +from litellm import encoding +from litellm import token_counter as litellm_token_counter +from litellm.types.utils import SelectTokenizerResponse +from tokenizers import Tokenizer + + def get_tokenizer_for_model(model: str) -> SelectTokenizerResponse: """Get the tokenizer for a model.""" try: