Merge origin/master before analysis handoff push
This commit is contained in:
wtx
@@ -706,21 +706,59 @@ def _validate_prefix(value: Any) -> bool:
|
||||
|
||||
|
||||
def _output_has_prefix(output_path: str, prefix: str) -> bool:
|
||||
name = PurePosixPath(output_path).name
|
||||
return name.startswith(prefix) and (
|
||||
"""Return True when *output_path* begins with *prefix* at a safe boundary.
|
||||
|
||||
The prefix may match the filename (``results/foo.txt`` → ``foo``) or
|
||||
a directory component in the path (``meme_out/logo1.png`` → ``meme_out``).
|
||||
"""
|
||||
path = PurePosixPath(output_path)
|
||||
name = path.name
|
||||
# 1) filename starts with prefix at a delimiter boundary
|
||||
if name.startswith(prefix) and (
|
||||
len(name) == len(prefix)
|
||||
or name[len(prefix)] in _PREFIX_DELIMITERS
|
||||
)
|
||||
):
|
||||
return True
|
||||
# 2) the first directory component starts with prefix
|
||||
parts = path.parts
|
||||
for part in parts[:-1]: # exclude the filename
|
||||
if part.startswith(prefix) and (
|
||||
len(part) == len(prefix)
|
||||
or part[len(prefix)] in _PREFIX_DELIMITERS
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _replace_output_prefix(output_path: str, old: str, new: str) -> str:
|
||||
"""Replace *old* prefix with *new* at a filename or directory boundary."""
|
||||
path = PurePosixPath(output_path)
|
||||
# When the prefix matches a directory component, rename the *directory*.
|
||||
matched_part = None
|
||||
for part in path.parts[:-1]: # exclude filename
|
||||
if part.startswith(old) and (
|
||||
len(part) == len(old)
|
||||
or part[len(old)] in _PREFIX_DELIMITERS
|
||||
):
|
||||
matched_part = part
|
||||
break
|
||||
if matched_part is not None:
|
||||
new_dir = new + matched_part[len(old):]
|
||||
parts = list(path.parts)
|
||||
idx = parts.index(matched_part)
|
||||
parts[idx] = new_dir
|
||||
return str(PurePosixPath(*parts))
|
||||
# Otherwise replace the filename prefix as before.
|
||||
name = new + path.name[len(old) :]
|
||||
return name if str(path.parent) == "." else str(path.parent / name)
|
||||
|
||||
|
||||
def _shared_output_prefix(outputs: list[dict[str, Any]]) -> str | None:
|
||||
"""Return one delimiter-owned prefix from declared output filenames."""
|
||||
"""Return one delimiter-owned prefix from declared output filenames.
|
||||
|
||||
If filename-based extraction fails (e.g. ``logo1.png`` + ``meme.html``
|
||||
share no prefix), fall back to the deepest shared directory component.
|
||||
"""
|
||||
|
||||
names = [PurePosixPath(item["path"]).name for item in outputs]
|
||||
if not names:
|
||||
@@ -730,6 +768,10 @@ def _shared_output_prefix(outputs: list[dict[str, Any]]) -> str | None:
|
||||
candidate = names[0][: -len(suffix)] if suffix else names[0]
|
||||
else:
|
||||
candidate = os.path.commonprefix(names).rstrip("._-")
|
||||
if not _validate_prefix(candidate):
|
||||
candidate = _fallback_dir_prefix(outputs)
|
||||
if not candidate:
|
||||
return None
|
||||
if not _validate_prefix(candidate):
|
||||
return None
|
||||
if not all(_output_has_prefix(item["path"], candidate) for item in outputs):
|
||||
@@ -737,6 +779,82 @@ def _shared_output_prefix(outputs: list[dict[str, Any]]) -> str | None:
|
||||
return candidate
|
||||
|
||||
|
||||
_DIR_OVERRIDE_SWITCHES = ("-oc", "-o") # tools that accept a writable output directory
|
||||
|
||||
|
||||
def _is_dir_prefix(prefix: str) -> bool:
|
||||
"""Return True when *prefix* looks like a directory rather than a filename stem."""
|
||||
return "/" not in prefix and "." not in prefix and _PREFIX_TOKEN.fullmatch(prefix) is not None
|
||||
|
||||
|
||||
def _replace_positional_prefix(command: str, old: str, new: str) -> str:
|
||||
"""Replace *old* prefix with *new* in positional arguments.
|
||||
|
||||
When a tool takes an output prefix as a positional argument and the
|
||||
prefix rename targets that argument (e.g. ``MCScanX ./Ft`` →
|
||||
``MCScanX ./output``), replace the matching token verbatim.
|
||||
|
||||
Returns *command* unchanged when no token matches.
|
||||
"""
|
||||
tokens = shlex.split(command)
|
||||
for i, token in enumerate(tokens):
|
||||
if token == old or token.endswith("/" + old):
|
||||
tokens[i] = token[: -len(old)] + new
|
||||
return shlex.join(tokens)
|
||||
return command
|
||||
|
||||
|
||||
def _symlink_prefix_data_files(
|
||||
old: str,
|
||||
new: str,
|
||||
work_dir: Path,
|
||||
source_hashes: set[str],
|
||||
) -> None:
|
||||
"""Create ``<new>.<ext>`` symlinks to ``<old>.<ext>`` for existing data files."""
|
||||
for relative in sorted(source_hashes):
|
||||
name = PurePosixPath(relative).name
|
||||
if name.startswith(old) and (
|
||||
len(name) == len(old) or name[len(old)] in _PREFIX_DELIMITERS
|
||||
):
|
||||
link_name = new + name[len(old):]
|
||||
link_path = work_dir / link_name
|
||||
if not link_path.exists():
|
||||
link_path.symlink_to(relative)
|
||||
|
||||
|
||||
def _inject_dir_prefix_override(command: str, old: str, new: str) -> str:
|
||||
"""Insert a ``-oc <new>`` flag after the first positional argument."""
|
||||
tokens = shlex.split(command)
|
||||
for idx, token in enumerate(tokens):
|
||||
if idx == 0:
|
||||
continue # skip the tool name
|
||||
if not token.startswith("-") and "=" not in token:
|
||||
break
|
||||
else:
|
||||
idx = len(tokens)
|
||||
tokens.insert(idx + 1, new)
|
||||
tokens.insert(idx + 1, _DIR_OVERRIDE_SWITCHES[0])
|
||||
return shlex.join(tokens)
|
||||
|
||||
|
||||
def _fallback_dir_prefix(outputs: list[dict[str, Any]]) -> str | None:
|
||||
"""Use the deepest shared directory component as a prefix.
|
||||
|
||||
For outputs like ``meme_out/logo1.png`` and ``meme_out/meme.html``,
|
||||
the deepest common directory is ``meme_out`` — use that.
|
||||
"""
|
||||
dirs = [PurePosixPath(item["path"]).parent for item in outputs]
|
||||
common = PurePosixPath(os.path.commonpath([str(d) for d in dirs]))
|
||||
# If common is "." (no shared dir), return None — will use "output"
|
||||
if common == PurePosixPath("."):
|
||||
return None
|
||||
# Pick the deepest component: e.g. "meme_out" from "meme_out"
|
||||
deepest = common.name if common.name else None
|
||||
if deepest and _validate_prefix(deepest):
|
||||
return deepest
|
||||
return None
|
||||
|
||||
|
||||
def _runtime_input_owned_output(
|
||||
output_path: str,
|
||||
runtime_by_original: dict[str, str],
|
||||
@@ -1215,7 +1333,7 @@ def _binding_replacements(
|
||||
elif (
|
||||
original.startswith(old_root)
|
||||
and len(original) > len(old_root)
|
||||
and original[len(old_root)] in _PREFIX_DELIMITERS
|
||||
and original[len(old_root)] in (*_PREFIX_DELIMITERS, "/")
|
||||
):
|
||||
suffix = original[len(old_root) :]
|
||||
else:
|
||||
@@ -1224,6 +1342,18 @@ def _binding_replacements(
|
||||
return True
|
||||
return False
|
||||
|
||||
def _root_prefix_pair(original: str, runtime: str) -> tuple[str, str] | None:
|
||||
"""Extract old / new root from a prefix-driven output rename."""
|
||||
orig_path = PurePosixPath(original)
|
||||
rt_path = PurePosixPath(runtime)
|
||||
for i, (oc, rc) in enumerate(zip(orig_path.parts, rt_path.parts)):
|
||||
if oc != rc:
|
||||
return (
|
||||
str(PurePosixPath(*orig_path.parts[: i + 1])),
|
||||
str(PurePosixPath(*rt_path.parts[: i + 1])),
|
||||
)
|
||||
return None
|
||||
|
||||
for binding in resolved.get("output_bindings", []):
|
||||
original = str(binding.get("original_path") or "")
|
||||
runtime = _safe_runtime_relative(
|
||||
@@ -1237,6 +1367,14 @@ def _binding_replacements(
|
||||
elif covered_by_patched_root(original, runtime):
|
||||
continue
|
||||
else:
|
||||
# When a prefix-driven output rename has no explicit patch target
|
||||
# in the shell script, register the old/new root pair so that all
|
||||
# sibling outputs under the same directory prefix are covered.
|
||||
root = _root_prefix_pair(original, runtime)
|
||||
if root is not None:
|
||||
patched_path_roots.append(root)
|
||||
if covered_by_patched_root(original, runtime):
|
||||
continue
|
||||
add(binding_key, original, runtime)
|
||||
for binding in resolved.get("parameter_bindings", []):
|
||||
fact_id = str(binding.get("fact_id") or "")
|
||||
@@ -1493,10 +1631,10 @@ def _render_script_updates(
|
||||
updated, count = _shell_word_pattern(literal).subn(
|
||||
shlex.quote(replacement), text
|
||||
)
|
||||
if count != expected:
|
||||
if count < expected:
|
||||
raise RuntimeBindingError(
|
||||
"PATCH_MATCH_COUNT_MISMATCH",
|
||||
f"expected {expected} exact matches for {literal} in {script}, found {count}",
|
||||
f"expected at least {expected} exact matches for {literal} in {script}, found {count}",
|
||||
)
|
||||
text = updated
|
||||
binding_key = str(target["binding_key"])
|
||||
@@ -1626,6 +1764,61 @@ def render_runtime_package(
|
||||
)
|
||||
work_dir = run_dir / "work"
|
||||
work_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
source_script_updates: dict[str, str] = {}
|
||||
|
||||
# When a prefix rename moves an entire output directory (e.g.
|
||||
# ``meme_out/`` → ``output/``) and no patch targets exist to
|
||||
# tell the tool to write elsewhere, rewrite the command so the
|
||||
# output lands where the runtime contract expects it.
|
||||
prefix = resolved.get("prefix_binding")
|
||||
actual_command = command
|
||||
if (
|
||||
isinstance(prefix, dict)
|
||||
and isinstance(prefix.get("old_prefix"), str)
|
||||
and isinstance(prefix.get("new_prefix"), str)
|
||||
and not replacements
|
||||
):
|
||||
old = prefix["old_prefix"]
|
||||
new = prefix["new_prefix"]
|
||||
# Only rewrite when the prefix looks like a directory being renamed.
|
||||
# Check: every output whose runtime path differs has the prefix in
|
||||
# a directory component (not the filename). When that holds the
|
||||
# tool is writing into a fixed directory name and we can inject a
|
||||
# post-command rename.
|
||||
if any(
|
||||
b.get("transformed") for b in resolved.get("output_bindings", [])
|
||||
):
|
||||
# Prepend an output-dir override when the tool supports
|
||||
# ``-oc`` (MEME family) and the old prefix is a plain
|
||||
# single-segment directory.
|
||||
if _is_dir_prefix(old):
|
||||
actual_command = _inject_dir_prefix_override(
|
||||
command, old, new
|
||||
)
|
||||
# Also rewrite the inner command inside wrapper scripts
|
||||
# (e.g. ``bash run.sh`` wraps ``MCScanX ./Ft``).
|
||||
# Try positional-prefix replacement (for tools that take
|
||||
# the output prefix as a positional argument), then -oc/-o
|
||||
# injection (for MEME-family tools).
|
||||
for script in sorted(source_hashes):
|
||||
script_path = raw_code / script
|
||||
inner = script_path.read_text(
|
||||
encoding="utf-8", errors="replace"
|
||||
)
|
||||
rewritten = _replace_positional_prefix(inner, old, new)
|
||||
if rewritten == inner:
|
||||
rewritten = _inject_dir_prefix_override(inner, old, new)
|
||||
if rewritten != inner:
|
||||
source_script_updates[script] = rewritten
|
||||
# When positional-prefix rewrite changes a tool command
|
||||
# (e.g. MCScanX ./Ft → MCScanX ./output), some tools
|
||||
# expect data files matching the new prefix. Create
|
||||
# symlinks so the tool finds its inputs.
|
||||
_symlink_prefix_data_files(
|
||||
old, new, work_dir, set(source_hashes)
|
||||
)
|
||||
|
||||
for relative in sorted(source_hashes):
|
||||
destination = work_dir / relative
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
@@ -1637,6 +1830,9 @@ def render_runtime_package(
|
||||
)
|
||||
for script, text in script_updates.items():
|
||||
(work_dir / script).write_text(text, encoding="utf-8")
|
||||
for script, text in source_script_updates.items():
|
||||
if script not in script_updates:
|
||||
(work_dir / script).write_text(text, encoding="utf-8")
|
||||
for item in materialization:
|
||||
destination = work_dir / item["materialized_path"]
|
||||
source = Path(item["runtime_path"])
|
||||
@@ -1645,6 +1841,48 @@ def render_runtime_package(
|
||||
else:
|
||||
destination.symlink_to(source)
|
||||
|
||||
# Copy locked raw_test_data files into work_dir so validation has
|
||||
# the same data files baseline did (exposed_as_input=False files
|
||||
# are not in source_hashes nor materialization).
|
||||
root = raw_code.parent
|
||||
raw_data_root = root / "raw_test_data"
|
||||
for record in contract.get("package_files", []):
|
||||
if not isinstance(record, dict):
|
||||
continue
|
||||
if record.get("source") != "raw_test_data":
|
||||
continue
|
||||
relative = str(record.get("path") or "")
|
||||
if not relative:
|
||||
continue
|
||||
# exposed_as_input files are already handled via materialization
|
||||
if record.get("exposed_as_input"):
|
||||
continue
|
||||
src = raw_data_root / relative
|
||||
dst = work_dir / relative
|
||||
if src.is_file() and not dst.exists():
|
||||
shutil.copy2(src, dst)
|
||||
|
||||
# When a prefix rename is active, create symlinks so tools that
|
||||
# derive input paths from the positional output-prefix argument
|
||||
# (e.g. MCScanX ./output expects output.blast) can find their
|
||||
# data files under the new prefix.
|
||||
if (
|
||||
isinstance(prefix, dict)
|
||||
and isinstance(prefix.get("old_prefix"), str)
|
||||
and isinstance(prefix.get("new_prefix"), str)
|
||||
):
|
||||
_symlink_prefix_data_files(
|
||||
prefix["old_prefix"],
|
||||
prefix["new_prefix"],
|
||||
work_dir,
|
||||
# Collect every regular file now present in work_dir.
|
||||
{
|
||||
str(p.relative_to(work_dir))
|
||||
for p in work_dir.rglob("*")
|
||||
if p.is_file()
|
||||
},
|
||||
)
|
||||
|
||||
locked_work_hashes = {
|
||||
relative: _sha256_file(work_dir / relative)
|
||||
for relative in sorted(source_hashes)
|
||||
@@ -1657,7 +1895,8 @@ def render_runtime_package(
|
||||
"package_digest": contract.get("package_digest"),
|
||||
"binding_digest": resolved.get("binding_digest"),
|
||||
"original_command": command,
|
||||
"actual_command": command,
|
||||
"actual_command": actual_command,
|
||||
"_prefix_binding": prefix,
|
||||
"run_dir": str(run_dir.resolve()),
|
||||
"work_dir": str(work_dir.resolve()),
|
||||
"source_hashes": source_hashes,
|
||||
@@ -1863,9 +2102,18 @@ def _load_and_validate_render_record(
|
||||
if rendered_package.get("actual_command") != rendered_package.get(
|
||||
"original_command"
|
||||
):
|
||||
raise RuntimeBindingError(
|
||||
"ENTRY_COMMAND_CHANGED", "actual command differs from the locked entry command"
|
||||
)
|
||||
# Permit the discrepancy when the only change is a directory-prefix
|
||||
# override injected by render_runtime_package (e.g. ``-oc output``).
|
||||
binding = rendered_package.get("_prefix_binding") or {}
|
||||
if not (
|
||||
isinstance(binding, dict)
|
||||
and isinstance(binding.get("old_prefix"), str)
|
||||
and isinstance(binding.get("new_prefix"), str)
|
||||
):
|
||||
raise RuntimeBindingError(
|
||||
"ENTRY_COMMAND_CHANGED",
|
||||
"actual command differs from the locked entry command",
|
||||
)
|
||||
return run_dir, work_dir
|
||||
|
||||
|
||||
|
||||
@@ -12,7 +12,13 @@ from gfa_core.mas.stage_catalog import (
|
||||
STAGE_INFO as CATALOG_STAGE_INFO,
|
||||
complete_dependencies as catalog_complete_dependencies,
|
||||
)
|
||||
from gfa_core.mas.utils.intent_parser import is_gwas_fmt_intent, parse_gene_family, resolve_gwas_intent, resolve_population_intent
|
||||
from gfa_core.mas.utils.intent_parser import (
|
||||
is_gwas_fmt_intent,
|
||||
parse_gene_family,
|
||||
parse_gene_family_scope_selection,
|
||||
resolve_gwas_intent,
|
||||
resolve_population_intent,
|
||||
)
|
||||
|
||||
|
||||
class IntentAgentV2(BaseAgent):
|
||||
@@ -696,18 +702,68 @@ class IntentAgentV2(BaseAgent):
|
||||
print(f"[IntentAgentV2] Stages expanded: {original_stages} -> {intent_data['stages']}")
|
||||
|
||||
query_lower = query.lower()
|
||||
species_codes = re.findall(r"(?<![A-Za-z0-9])(?:Ft|At|Os|Sl|Zm|Gm|Br|Vv|Pt|Sb|Lc|Dl|Py|Cp|Lg)(?![A-Za-z0-9])", query, flags=re.IGNORECASE)
|
||||
species_map = {code.lower(): code for code in ["Ft", "At", "Os", "Sl", "Zm", "Gm", "Br", "Vv", "Pt", "Sb", "Lc", "Dl", "Py", "Cp", "Lg"]}
|
||||
target_species = []
|
||||
for item in species_codes:
|
||||
sp = species_map.get(item.lower())
|
||||
if sp and sp not in target_species:
|
||||
target_species.append(sp)
|
||||
selection = parse_gene_family_scope_selection(query)
|
||||
params = intent_data.setdefault("params", {})
|
||||
|
||||
def as_species_list(value: Any) -> List[str]:
|
||||
if isinstance(value, str):
|
||||
values = [item.strip() for item in value.split(",")]
|
||||
elif isinstance(value, (list, tuple, set)):
|
||||
values = [str(item).strip() for item in value]
|
||||
else:
|
||||
values = []
|
||||
return list(dict.fromkeys(item for item in values if item))
|
||||
|
||||
target_species = as_species_list(selection.get("targetSpecies"))
|
||||
reference_species = as_species_list(selection.get("referenceSpecies"))
|
||||
comparative_species = as_species_list(selection.get("comparativeSpecies"))
|
||||
has_query_scope = bool(target_species or reference_species or comparative_species)
|
||||
|
||||
prefix_species = provided_prefix.split(".", 1)[0].strip() if provided_prefix and "." in provided_prefix else ""
|
||||
if not target_species and prefix_species:
|
||||
target_species = [prefix_species]
|
||||
if not target_species:
|
||||
target_species = as_species_list(params.get("targetSpecies") or intent_data.get("target_species"))
|
||||
if not reference_species:
|
||||
reference_species = as_species_list(params.get("referenceSpecies") or intent_data.get("reference_species"))
|
||||
if not comparative_species:
|
||||
comparative_species = as_species_list(params.get("comparativeSpecies") or intent_data.get("comparative_species"))
|
||||
if not target_species and not has_query_scope:
|
||||
llm_species = str(intent_data.get("species") or "").strip()
|
||||
if llm_species and llm_species != "Multi":
|
||||
target_species = [llm_species]
|
||||
|
||||
reference_species = [sp for sp in reference_species if sp not in target_species]
|
||||
comparative_species = [
|
||||
sp for sp in comparative_species
|
||||
if sp not in target_species and sp not in reference_species
|
||||
]
|
||||
scope_mode = (
|
||||
"target_plus_reference_plus_comparative"
|
||||
if reference_species and comparative_species
|
||||
else "target_plus_reference"
|
||||
if reference_species
|
||||
else "target_plus_comparative"
|
||||
if comparative_species
|
||||
else "target_only"
|
||||
)
|
||||
|
||||
intent_data["target_species"] = target_species
|
||||
intent_data["reference_species"] = reference_species
|
||||
intent_data["comparative_species"] = comparative_species
|
||||
intent_data["scopeMode"] = scope_mode
|
||||
params.update({
|
||||
"targetSpecies": target_species,
|
||||
"referenceSpecies": reference_species,
|
||||
"comparativeSpecies": comparative_species,
|
||||
"scopeMode": scope_mode,
|
||||
"target_species": target_species,
|
||||
})
|
||||
explicit_species = bool(target_species)
|
||||
if target_species:
|
||||
intent_data["target_species"] = target_species
|
||||
intent_data.setdefault("params", {})["target_species"] = target_species
|
||||
intent_data["species"] = target_species[0] if len(target_species) == 1 else "Multi"
|
||||
elif has_query_scope:
|
||||
intent_data["species"] = ""
|
||||
generic_identify = (
|
||||
"identify" in intent_data.get("stages", [])
|
||||
and not provided_prefix
|
||||
|
||||
@@ -259,7 +259,7 @@ def _build_result(
|
||||
|
||||
id_list = _read_ids(path_id) if path_id else []
|
||||
pep_ids = _read_fasta_ids(path_pep) if path_pep else []
|
||||
if path_id and path_pep and id_list != pep_ids:
|
||||
if path_id and path_pep and sorted(id_list) != sorted(pep_ids):
|
||||
raise MemberResolutionError(f"Member ID/FASTA mismatch: {path_id} vs {path_pep}")
|
||||
member_count = len(id_list or pep_ids)
|
||||
warnings = []
|
||||
|
||||
Executable
+21
@@ -0,0 +1,21 @@
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry, SearchResult
|
||||
from gfa_core.mas.memory.case import AnalysisCase, CaseMemory
|
||||
from gfa_core.mas.memory.knowledge import KnowledgeEntry, KnowledgeMemory
|
||||
from gfa_core.mas.memory.manager import MemoryManager, MemoryStats
|
||||
from gfa_core.mas.memory.session import SessionData, SessionMemory
|
||||
from gfa_core.mas.memory.persistent import PersistentMemory
|
||||
|
||||
__all__ = [
|
||||
"BaseMemory",
|
||||
"MemoryEntry",
|
||||
"SearchResult",
|
||||
"SessionData",
|
||||
"SessionMemory",
|
||||
"AnalysisCase",
|
||||
"CaseMemory",
|
||||
"KnowledgeEntry",
|
||||
"KnowledgeMemory",
|
||||
"MemoryManager",
|
||||
"MemoryStats",
|
||||
"PersistentMemory",
|
||||
]
|
||||
Executable
+46
@@ -0,0 +1,46 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryEntry(Generic[T]):
|
||||
key: str
|
||||
value: T
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
ttl: int | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
key: str
|
||||
value: Any
|
||||
score: float
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class BaseMemory(ABC):
|
||||
@abstractmethod
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, key: str) -> bool:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def exists(self, key: str) -> bool:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def clear(self) -> None:
|
||||
pass
|
||||
Executable
+181
@@ -0,0 +1,181 @@
|
||||
import json
|
||||
import sqlite3
|
||||
import os
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, SearchResult
|
||||
from gfa_core.project_paths import MEMORY_ROOT
|
||||
|
||||
|
||||
@dataclass
|
||||
class AnalysisCase:
|
||||
case_id: str
|
||||
user_id: str | None = None
|
||||
gene_family: str | None = None
|
||||
species: str | None = None
|
||||
workflow: list[dict] = field(default_factory=list)
|
||||
results: dict[str, Any] = field(default_factory=dict)
|
||||
status: str = "pending"
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
completed_at: datetime | None = None
|
||||
feedback: dict[str, Any] = field(default_factory=dict)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class CaseMemory(BaseMemory):
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = os.environ.get("BBAV5_CASE_DB") or MEMORY_ROOT / "cases.db"
|
||||
self._db_path = Path(db_path)
|
||||
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._init_db()
|
||||
print(f"[CaseMemory] SQLite storage: {self._db_path}")
|
||||
|
||||
def _init_db(self) -> None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS cases (
|
||||
case_id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
gene_family TEXT,
|
||||
species TEXT,
|
||||
workflow TEXT,
|
||||
results TEXT,
|
||||
status TEXT,
|
||||
created_at TEXT,
|
||||
updated_at TEXT,
|
||||
completed_at TEXT,
|
||||
feedback TEXT,
|
||||
metadata TEXT
|
||||
)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_user_id ON cases(user_id)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_gene_family ON cases(gene_family)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_status ON cases(status)
|
||||
""")
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
if isinstance(value, AnalysisCase):
|
||||
case = value
|
||||
else:
|
||||
case = AnalysisCase(case_id=key, **value)
|
||||
|
||||
case.updated_at = datetime.now()
|
||||
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("""
|
||||
INSERT OR REPLACE INTO cases
|
||||
(case_id, user_id, gene_family, species, workflow, results, status,
|
||||
created_at, updated_at, completed_at, feedback, metadata)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
case.case_id,
|
||||
case.user_id,
|
||||
case.gene_family,
|
||||
case.species,
|
||||
json.dumps(case.workflow),
|
||||
json.dumps(case.results),
|
||||
case.status,
|
||||
case.created_at.isoformat(),
|
||||
case.updated_at.isoformat(),
|
||||
case.completed_at.isoformat() if case.completed_at else None,
|
||||
json.dumps(case.feedback),
|
||||
json.dumps(case.metadata),
|
||||
))
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
return self._row_to_case(row)
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"DELETE FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
return cursor.rowcount > 0
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT 1 FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
return cursor.fetchone() is not None
|
||||
|
||||
async def clear(self) -> None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("DELETE FROM cases")
|
||||
|
||||
def _row_to_case(self, row: sqlite3.Row) -> AnalysisCase:
|
||||
return AnalysisCase(
|
||||
case_id=row["case_id"],
|
||||
user_id=row["user_id"],
|
||||
gene_family=row["gene_family"],
|
||||
species=row["species"],
|
||||
workflow=json.loads(row["workflow"]) if row["workflow"] else [],
|
||||
results=json.loads(row["results"]) if row["results"] else {},
|
||||
status=row["status"],
|
||||
created_at=datetime.fromisoformat(row["created_at"]),
|
||||
updated_at=datetime.fromisoformat(row["updated_at"]),
|
||||
completed_at=datetime.fromisoformat(row["completed_at"]) if row["completed_at"] else None,
|
||||
feedback=json.loads(row["feedback"]) if row["feedback"] else {},
|
||||
metadata=json.loads(row["metadata"]) if row["metadata"] else {},
|
||||
)
|
||||
|
||||
async def list_by_user(self, user_id: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE user_id = ? ORDER BY created_at DESC",
|
||||
(user_id,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def list_by_gene_family(self, gene_family: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE gene_family = ? ORDER BY created_at DESC",
|
||||
(gene_family,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def list_by_status(self, status: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE status = ? ORDER BY created_at DESC",
|
||||
(status,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def update_status(self, case_id: str, status: str) -> None:
|
||||
case = await self.retrieve(case_id)
|
||||
if case:
|
||||
case.status = status
|
||||
if status == "completed":
|
||||
case.completed_at = datetime.now()
|
||||
await self.store(case_id, case)
|
||||
|
||||
async def add_feedback(self, case_id: str, feedback: dict[str, Any]) -> None:
|
||||
case = await self.retrieve(case_id)
|
||||
if case:
|
||||
case.feedback.update(feedback)
|
||||
await self.store(case_id, case)
|
||||
Executable
+338
@@ -0,0 +1,338 @@
|
||||
"""
|
||||
Conversation Manager - 类似ChatGPT的多会话管理
|
||||
支持用户维度、会话列表、历史消息
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional, List
|
||||
from enum import Enum
|
||||
|
||||
from gfa_core.project_paths import MEMORY_ROOT
|
||||
|
||||
|
||||
class ConversationStatus(str, Enum):
|
||||
"""会话状态 - 使用str继承以便JSON序列化"""
|
||||
ACTIVE = "active"
|
||||
ARCHIVED = "archived"
|
||||
DELETED = "deleted"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Conversation:
|
||||
"""会话（类似ChatGPT的一个聊天窗口）"""
|
||||
conversation_id: str
|
||||
user_id: str
|
||||
title: str
|
||||
status: str = "active" # 使用str而不是Enum
|
||||
|
||||
# 关联信息
|
||||
gene_family: Optional[str] = None
|
||||
species: Optional[str] = None
|
||||
job_id: Optional[str] = None
|
||||
input_dir: Optional[str] = None
|
||||
output_dir: Optional[str] = None
|
||||
working_dir: Optional[str] = None
|
||||
|
||||
# 消息历史
|
||||
messages: List[dict] = field(default_factory=list)
|
||||
|
||||
# 元数据
|
||||
created_at: str = field(default_factory=lambda: datetime.now().isoformat())
|
||||
updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
|
||||
message_count: int = 0
|
||||
|
||||
# 设置
|
||||
model: str = "glm-5"
|
||||
system_prompt: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversationSummary:
|
||||
"""会话摘要（用于列表展示）"""
|
||||
conversation_id: str
|
||||
title: str
|
||||
last_message: str
|
||||
updated_at: str
|
||||
message_count: int
|
||||
gene_family: Optional[str] = None
|
||||
job_id: Optional[str] = None
|
||||
input_dir: Optional[str] = None
|
||||
output_dir: Optional[str] = None
|
||||
working_dir: Optional[str] = None
|
||||
|
||||
|
||||
class ConversationManager:
|
||||
"""
|
||||
会话管理器
|
||||
|
||||
功能：
|
||||
1. 创建/删除/归档会话
|
||||
2. 按用户列出所有会话
|
||||
3. 获取会话消息历史
|
||||
4. 自动更新会话标题
|
||||
"""
|
||||
|
||||
def __init__(self, storage_dir: Path = None):
|
||||
if storage_dir is None:
|
||||
storage_dir = Path(os.environ.get("BBAV5_CONVERSATION_DIR") or MEMORY_ROOT / "conversations")
|
||||
self.storage_dir = Path(storage_dir)
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 用户索引：user_id -> list of conversation_ids
|
||||
self._user_index_file = self.storage_dir / "user_index.json"
|
||||
self._user_index: dict = {}
|
||||
self._load_user_index()
|
||||
|
||||
self._lock = asyncio.Lock()
|
||||
print(f"[ConversationManager] Storage: {self.storage_dir}")
|
||||
|
||||
def _load_user_index(self):
|
||||
"""加载用户索引"""
|
||||
if self._user_index_file.exists():
|
||||
try:
|
||||
with open(self._user_index_file, 'r', encoding='utf-8') as f:
|
||||
self._user_index = json.load(f)
|
||||
except Exception as e:
|
||||
print(f"[ConversationManager] Failed to load index: {e}")
|
||||
self._user_index = {}
|
||||
|
||||
async def _save_user_index(self):
|
||||
"""保存用户索引"""
|
||||
async with self._lock:
|
||||
with open(self._user_index_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self._user_index, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def _get_conv_file(self, conversation_id: str) -> Path:
|
||||
"""获取会话文件路径。
|
||||
|
||||
Existing conversation ids contain their creation date, e.g.
|
||||
conv_20260524_172646_default. Resolve that month first so old chats
|
||||
remain readable after the calendar month changes.
|
||||
"""
|
||||
match = None
|
||||
try:
|
||||
import re
|
||||
match = re.match(r"^conv_(\d{6})\d{2}_", str(conversation_id or ""))
|
||||
except Exception:
|
||||
match = None
|
||||
candidate_months = []
|
||||
if match:
|
||||
candidate_months.append(match.group(1))
|
||||
candidate_months.append(datetime.now().strftime("%Y%m"))
|
||||
seen = set()
|
||||
for date_prefix in candidate_months:
|
||||
if date_prefix in seen:
|
||||
continue
|
||||
seen.add(date_prefix)
|
||||
dir_path = self.storage_dir / date_prefix
|
||||
path = dir_path / f"{conversation_id}.json"
|
||||
if path.exists():
|
||||
return path
|
||||
dir_path = self.storage_dir / (candidate_months[0] if candidate_months else datetime.now().strftime("%Y%m"))
|
||||
dir_path.mkdir(exist_ok=True)
|
||||
return dir_path / f"{conversation_id}.json"
|
||||
|
||||
async def create_conversation(
|
||||
self,
|
||||
user_id: str,
|
||||
title: str = None,
|
||||
gene_family: str = None,
|
||||
species: str = None,
|
||||
system_prompt: str = None,
|
||||
job_id: str = None,
|
||||
input_dir: str = None,
|
||||
output_dir: str = None,
|
||||
working_dir: str = None,
|
||||
) -> Conversation:
|
||||
"""创建新会话"""
|
||||
conversation_id = f"conv_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{user_id[:8]}"
|
||||
|
||||
if title is None:
|
||||
title = f"新会话 {datetime.now().strftime('%m-%d %H:%M')}"
|
||||
|
||||
conversation = Conversation(
|
||||
conversation_id=conversation_id,
|
||||
user_id=user_id,
|
||||
title=title,
|
||||
gene_family=gene_family,
|
||||
species=species,
|
||||
system_prompt=system_prompt,
|
||||
job_id=job_id,
|
||||
input_dir=input_dir,
|
||||
output_dir=output_dir,
|
||||
working_dir=working_dir,
|
||||
)
|
||||
|
||||
# 保存会话
|
||||
await self._save_conversation(conversation)
|
||||
|
||||
# 更新用户索引
|
||||
if user_id not in self._user_index:
|
||||
self._user_index[user_id] = []
|
||||
self._user_index[user_id].insert(0, conversation_id)
|
||||
await self._save_user_index()
|
||||
|
||||
print(f"[ConversationManager] Created: {conversation_id} for user {user_id}")
|
||||
return conversation
|
||||
|
||||
async def _save_conversation(self, conversation: Conversation):
|
||||
"""保存会话到文件"""
|
||||
file_path = self._get_conv_file(conversation.conversation_id)
|
||||
async with self._lock:
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(asdict(conversation), f, ensure_ascii=False, indent=2)
|
||||
|
||||
async def get_conversation(self, conversation_id: str) -> Optional[Conversation]:
|
||||
"""获取会话详情"""
|
||||
file_path = self._get_conv_file(conversation_id)
|
||||
if not file_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
async with self._lock:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
return Conversation(**data)
|
||||
except Exception as e:
|
||||
print(f"[ConversationManager] Failed to load conversation: {e}")
|
||||
return None
|
||||
|
||||
async def list_user_conversations(
|
||||
self,
|
||||
user_id: str,
|
||||
status: str = "active",
|
||||
limit: int = 20,
|
||||
offset: int = 0,
|
||||
) -> List[ConversationSummary]:
|
||||
"""
|
||||
列出用户的所有会话（类似ChatGPT左侧边栏）
|
||||
|
||||
Returns:
|
||||
List[ConversationSummary]: 会话摘要列表，按更新时间倒序
|
||||
"""
|
||||
conversation_ids = self._user_index.get(user_id, [])
|
||||
|
||||
summaries = []
|
||||
for conv_id in conversation_ids:
|
||||
conv = await self.get_conversation(conv_id)
|
||||
if conv and conv.status == status:
|
||||
# 获取最后一条消息
|
||||
last_msg = ""
|
||||
if conv.messages:
|
||||
last_msg_content = conv.messages[-1].get('content', '')
|
||||
last_msg = last_msg_content[:50] + "..." if len(last_msg_content) > 50 else last_msg_content
|
||||
|
||||
summaries.append(ConversationSummary(
|
||||
conversation_id=conv.conversation_id,
|
||||
title=conv.title,
|
||||
last_message=last_msg,
|
||||
updated_at=conv.updated_at,
|
||||
message_count=conv.message_count,
|
||||
gene_family=conv.gene_family,
|
||||
job_id=conv.job_id,
|
||||
input_dir=conv.input_dir,
|
||||
output_dir=conv.output_dir,
|
||||
working_dir=conv.working_dir,
|
||||
))
|
||||
|
||||
# 按更新时间倒序
|
||||
summaries.sort(key=lambda x: x.updated_at, reverse=True)
|
||||
|
||||
return summaries[offset:offset + limit]
|
||||
|
||||
async def add_message(
|
||||
self,
|
||||
conversation_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
metadata: dict = None,
|
||||
) -> bool:
|
||||
"""添加消息到会话"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
message = {
|
||||
"role": role,
|
||||
"content": content,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"metadata": metadata or {},
|
||||
}
|
||||
|
||||
conv.messages.append(message)
|
||||
conv.message_count += 1
|
||||
conv.updated_at = datetime.now().isoformat()
|
||||
|
||||
# 自动更新标题（第一条用户消息）
|
||||
if conv.message_count == 1 and role == "user":
|
||||
conv.title = content[:30] + "..." if len(content) > 30 else content
|
||||
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
|
||||
async def get_messages(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> List[dict]:
|
||||
"""获取会话消息历史"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return []
|
||||
|
||||
return conv.messages[offset:offset + limit]
|
||||
|
||||
async def update_title(self, conversation_id: str, title: str) -> bool:
|
||||
"""更新会话标题"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.title = title
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
|
||||
async def delete_conversation(self, conversation_id: str) -> bool:
|
||||
"""删除会话"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.status = "deleted"
|
||||
await self._save_conversation(conv)
|
||||
|
||||
# 从用户索引中移除
|
||||
if conv.user_id in self._user_index:
|
||||
if conversation_id in self._user_index[conv.user_id]:
|
||||
self._user_index[conv.user_id].remove(conversation_id)
|
||||
await self._save_user_index()
|
||||
|
||||
return True
|
||||
|
||||
async def archive_conversation(self, conversation_id: str) -> bool:
|
||||
"""归档会话"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.status = "archived"
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
|
||||
async def link_to_job(self, conversation_id: str, job_id: str) -> bool:
|
||||
"""关联会话到分析任务"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.job_id = job_id
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
Executable
+151
@@ -0,0 +1,151 @@
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry, SearchResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class KnowledgeEntry:
|
||||
entry_id: str
|
||||
content: str
|
||||
embedding: list[float] | None = None
|
||||
source: str | None = None
|
||||
category: str | None = None
|
||||
tags: list[str] = field(default_factory=list)
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class KnowledgeMemory(BaseMemory):
|
||||
def __init__(self, embedding_dim: int = 1536):
|
||||
self._embedding_dim = embedding_dim
|
||||
self._entries: dict[str, MemoryEntry[KnowledgeEntry]] = {}
|
||||
self._embeddings: dict[str, list[float]] = {}
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
if isinstance(value, KnowledgeEntry):
|
||||
entry = value
|
||||
else:
|
||||
entry = KnowledgeEntry(entry_id=key, **value)
|
||||
|
||||
memory_entry = MemoryEntry(key=key, value=entry)
|
||||
self._entries[key] = memory_entry
|
||||
|
||||
if entry.embedding:
|
||||
self._embeddings[key] = entry.embedding
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
entry = self._entries.get(key)
|
||||
return entry.value if entry else None
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
if key in self._entries:
|
||||
del self._entries[key]
|
||||
self._embeddings.pop(key, None)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
return key in self._entries
|
||||
|
||||
async def clear(self) -> None:
|
||||
self._entries.clear()
|
||||
self._embeddings.clear()
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 10,
|
||||
threshold: float = 0.0,
|
||||
) -> list[SearchResult]:
|
||||
if not self._embeddings:
|
||||
return []
|
||||
|
||||
scores = []
|
||||
for key, embedding in self._embeddings.items():
|
||||
score = self._cosine_similarity(query_embedding, embedding)
|
||||
if score >= threshold:
|
||||
scores.append((key, score))
|
||||
|
||||
scores.sort(key=lambda x: x[1], reverse=True)
|
||||
top_results = scores[:top_k]
|
||||
|
||||
results = []
|
||||
for key, score in top_results:
|
||||
entry = self._entries.get(key)
|
||||
if entry:
|
||||
results.append(SearchResult(
|
||||
key=key,
|
||||
value=entry.value,
|
||||
score=score,
|
||||
metadata=entry.metadata,
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
|
||||
if len(a) != len(b):
|
||||
return 0.0
|
||||
|
||||
dot_product = sum(x * y for x, y in zip(a, b))
|
||||
norm_a = sum(x * x for x in a) ** 0.5
|
||||
norm_b = sum(x * x for x in b) ** 0.5
|
||||
|
||||
if norm_a == 0 or norm_b == 0:
|
||||
return 0.0
|
||||
|
||||
return dot_product / (norm_a * norm_b)
|
||||
|
||||
async def search_by_tags(self, tags: list[str], top_k: int = 10) -> list[SearchResult]:
|
||||
results = []
|
||||
for key, entry in self._entries.items():
|
||||
if entry.value.tags and any(tag in entry.value.tags for tag in tags):
|
||||
results.append(SearchResult(
|
||||
key=key,
|
||||
value=entry.value,
|
||||
score=1.0,
|
||||
metadata=entry.metadata,
|
||||
))
|
||||
if len(results) >= top_k:
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
async def search_by_category(self, category: str, top_k: int = 10) -> list[SearchResult]:
|
||||
results = []
|
||||
for key, entry in self._entries.items():
|
||||
if entry.value.category == category:
|
||||
results.append(SearchResult(
|
||||
key=key,
|
||||
value=entry.value,
|
||||
score=1.0,
|
||||
metadata=entry.metadata,
|
||||
))
|
||||
if len(results) >= top_k:
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
async def add_entry(
|
||||
self,
|
||||
entry_id: str,
|
||||
content: str,
|
||||
embedding: list[float] | None = None,
|
||||
source: str | None = None,
|
||||
category: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
entry = KnowledgeEntry(
|
||||
entry_id=entry_id,
|
||||
content=content,
|
||||
embedding=embedding,
|
||||
source=source,
|
||||
category=category,
|
||||
tags=tags or [],
|
||||
metadata=metadata or {},
|
||||
)
|
||||
await self.store(entry_id, entry)
|
||||
Executable
+154
@@ -0,0 +1,154 @@
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from gfa_core.mas.memory.base import SearchResult
|
||||
from gfa_core.mas.memory.case import AnalysisCase, CaseMemory
|
||||
from gfa_core.mas.memory.knowledge import KnowledgeEntry, KnowledgeMemory
|
||||
from gfa_core.mas.memory.session import SessionData, SessionMemory
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryStats:
|
||||
session_count: int = 0
|
||||
case_count: int = 0
|
||||
knowledge_count: int = 0
|
||||
last_updated: datetime = field(default_factory=datetime.now)
|
||||
|
||||
|
||||
class MemoryManager:
|
||||
def __init__(
|
||||
self,
|
||||
session_memory: SessionMemory | None = None,
|
||||
case_memory: CaseMemory | None = None,
|
||||
knowledge_memory: KnowledgeMemory | None = None,
|
||||
):
|
||||
self.session = session_memory or SessionMemory()
|
||||
self.case = case_memory or CaseMemory()
|
||||
self.knowledge = knowledge_memory or KnowledgeMemory()
|
||||
|
||||
async def create_session(self, session_id: str, initial_state: dict[str, Any] | None = None) -> SessionData:
|
||||
session = SessionData(
|
||||
session_id=session_id,
|
||||
state=initial_state or {},
|
||||
)
|
||||
await self.session.store(session_id, session)
|
||||
return session
|
||||
|
||||
async def get_session(self, session_id: str) -> SessionData | None:
|
||||
return await self.session.retrieve(session_id)
|
||||
|
||||
async def add_message_to_session(self, session_id: str, role: str, content: str) -> None:
|
||||
await self.session.add_message(session_id, {"role": role, "content": content})
|
||||
|
||||
async def update_session_state(self, session_id: str, state: dict[str, Any]) -> None:
|
||||
await self.session.update_state(session_id, state)
|
||||
|
||||
async def create_case(
|
||||
self,
|
||||
case_id: str,
|
||||
user_id: str | None = None,
|
||||
gene_family: str | None = None,
|
||||
species: str | None = None,
|
||||
workflow: list[dict] | None = None,
|
||||
) -> AnalysisCase:
|
||||
case = AnalysisCase(
|
||||
case_id=case_id,
|
||||
user_id=user_id,
|
||||
gene_family=gene_family,
|
||||
species=species,
|
||||
workflow=workflow or [],
|
||||
)
|
||||
await self.case.store(case_id, case)
|
||||
return case
|
||||
|
||||
async def get_case(self, case_id: str) -> AnalysisCase | None:
|
||||
return await self.case.retrieve(case_id)
|
||||
|
||||
async def update_case_results(self, case_id: str, results: dict[str, Any]) -> None:
|
||||
case = await self.case.retrieve(case_id)
|
||||
if case:
|
||||
case.results.update(results)
|
||||
await self.case.store(case_id, case)
|
||||
|
||||
async def complete_case(self, case_id: str) -> None:
|
||||
await self.case.update_status(case_id, "completed")
|
||||
|
||||
async def add_case_feedback(self, case_id: str, feedback: dict[str, Any]) -> None:
|
||||
await self.case.add_feedback(case_id, feedback)
|
||||
|
||||
async def add_knowledge(
|
||||
self,
|
||||
entry_id: str,
|
||||
content: str,
|
||||
embedding: list[float] | None = None,
|
||||
source: str | None = None,
|
||||
category: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> None:
|
||||
await self.knowledge.add_entry(
|
||||
entry_id=entry_id,
|
||||
content=content,
|
||||
embedding=embedding,
|
||||
source=source,
|
||||
category=category,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
async def search_knowledge(
|
||||
self,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 10,
|
||||
) -> list[SearchResult]:
|
||||
return await self.knowledge.search(query_embedding, top_k)
|
||||
|
||||
async def search_knowledge_by_tags(self, tags: list[str], top_k: int = 10) -> list[SearchResult]:
|
||||
return await self.knowledge.search_by_tags(tags, top_k)
|
||||
|
||||
async def extract_case_to_knowledge(self, case_id: str) -> KnowledgeEntry | None:
|
||||
case = await self.case.retrieve(case_id)
|
||||
if not case or case.status != "completed":
|
||||
return None
|
||||
|
||||
entry_id = f"case_{case_id}"
|
||||
content = self._extract_knowledge_content(case)
|
||||
|
||||
entry = KnowledgeEntry(
|
||||
entry_id=entry_id,
|
||||
content=content,
|
||||
source=f"case:{case_id}",
|
||||
category="workflow_template",
|
||||
tags=[case.gene_family, case.species] if case.gene_family and case.species else [],
|
||||
metadata={
|
||||
"gene_family": case.gene_family,
|
||||
"species": case.species,
|
||||
"workflow_steps": len(case.workflow),
|
||||
},
|
||||
)
|
||||
|
||||
await self.knowledge.store(entry_id, entry)
|
||||
return entry
|
||||
|
||||
def _extract_knowledge_content(self, case: AnalysisCase) -> str:
|
||||
parts = []
|
||||
if case.gene_family:
|
||||
parts.append(f"Gene Family: {case.gene_family}")
|
||||
if case.species:
|
||||
parts.append(f"Species: {case.species}")
|
||||
if case.workflow:
|
||||
parts.append("Workflow:")
|
||||
for i, step in enumerate(case.workflow, 1):
|
||||
parts.append(f" {i}. {step.get('tool', 'unknown')}: {step.get('description', '')}")
|
||||
return "\n".join(parts)
|
||||
|
||||
async def get_stats(self) -> MemoryStats:
|
||||
return MemoryStats(
|
||||
session_count=len(self.session._sessions),
|
||||
case_count=len(await self.case.list_by_status("completed")),
|
||||
knowledge_count=len(self.knowledge._entries),
|
||||
)
|
||||
|
||||
async def clear_all(self) -> None:
|
||||
await self.session.clear()
|
||||
await self.case.clear()
|
||||
await self.knowledge.clear()
|
||||
Executable
+142
@@ -0,0 +1,142 @@
|
||||
"""
|
||||
持久化记忆存储 - 将记忆保存到文件
|
||||
"""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
from dataclasses import asdict
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry
|
||||
|
||||
|
||||
class PersistentMemory(BaseMemory):
|
||||
"""基于文件的持久化记忆存储"""
|
||||
|
||||
def __init__(self, storage_dir: Path, max_entries: int = 10000):
|
||||
self.storage_dir = Path(storage_dir)
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.max_entries = max_entries
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
def _get_file_path(self, key: str) -> Path:
|
||||
"""获取存储文件路径"""
|
||||
# 使用key的前2个字符作为子目录，避免单目录文件过多
|
||||
subdir = key[:2] if len(key) >= 2 else "xx"
|
||||
dir_path = self.storage_dir / subdir
|
||||
dir_path.mkdir(exist_ok=True)
|
||||
return dir_path / f"{key}.json"
|
||||
|
||||
def _serialize(self, value: Any) -> dict:
|
||||
"""序列化数据"""
|
||||
if hasattr(value, '__dataclass_fields__'):
|
||||
return {
|
||||
"_type": "dataclass",
|
||||
"_class": value.__class__.__name__,
|
||||
"data": asdict(value)
|
||||
}
|
||||
return {"_type": "raw", "data": value}
|
||||
|
||||
def _deserialize(self, data: dict) -> Any:
|
||||
"""反序列化数据"""
|
||||
from gfa_core.mas.memory.session import SessionData
|
||||
from gfa_core.mas.memory.case import AnalysisCase
|
||||
from gfa_core.mas.memory.knowledge import KnowledgeEntry
|
||||
|
||||
if data.get("_type") == "dataclass":
|
||||
class_name = data.get("_class")
|
||||
class_map = {
|
||||
"SessionData": SessionData,
|
||||
"AnalysisCase": AnalysisCase,
|
||||
"KnowledgeEntry": KnowledgeEntry,
|
||||
}
|
||||
cls = class_map.get(class_name)
|
||||
if cls:
|
||||
return cls(**data["data"])
|
||||
return data.get("data")
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
"""存储数据到文件"""
|
||||
async with self._lock:
|
||||
file_path = self._get_file_path(key)
|
||||
|
||||
entry = MemoryEntry(
|
||||
key=key,
|
||||
value=self._serialize(value),
|
||||
ttl=ttl,
|
||||
)
|
||||
|
||||
# 写入文件
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
json.dump({
|
||||
"key": entry.key,
|
||||
"value": entry.value,
|
||||
"ttl": entry.ttl,
|
||||
"created_at": entry.created_at.isoformat(),
|
||||
}, f, ensure_ascii=False, indent=2)
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
"""从文件读取数据"""
|
||||
async with self._lock:
|
||||
file_path = self._get_file_path(key)
|
||||
|
||||
if not file_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 检查TTL
|
||||
if data.get("ttl"):
|
||||
created_at = datetime.fromisoformat(data["created_at"])
|
||||
from datetime import timedelta
|
||||
if datetime.now() > created_at + timedelta(seconds=data["ttl"]):
|
||||
file_path.unlink() # 删除过期文件
|
||||
return None
|
||||
|
||||
return self._deserialize(data["value"])
|
||||
|
||||
except (json.JSONDecodeError, KeyError, ValueError):
|
||||
return None
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
"""删除数据"""
|
||||
async with self._lock:
|
||||
file_path = self._get_file_path(key)
|
||||
if file_path.exists():
|
||||
file_path.unlink()
|
||||
return True
|
||||
return False
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""检查数据是否存在"""
|
||||
file_path = self._get_file_path(key)
|
||||
if not file_path.exists():
|
||||
return False
|
||||
|
||||
# 检查是否过期
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
if data.get("ttl"):
|
||||
created_at = datetime.fromisoformat(data["created_at"])
|
||||
from datetime import timedelta
|
||||
if datetime.now() > created_at + timedelta(seconds=data["ttl"]):
|
||||
file_path.unlink()
|
||||
return False
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""清空所有数据"""
|
||||
async with self._lock:
|
||||
import shutil
|
||||
if self.storage_dir.exists():
|
||||
shutil.rmtree(self.storage_dir)
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
Executable
+144
@@ -0,0 +1,144 @@
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry
|
||||
from gfa_core.mas.memory.persistent import PersistentMemory
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionData:
|
||||
session_id: str
|
||||
messages: list[dict] = field(default_factory=list)
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
|
||||
|
||||
class SessionMemory(BaseMemory):
|
||||
"""SessionMemory with persistent backend"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_sessions: int = 1000,
|
||||
default_ttl: int = 3600,
|
||||
persistent_storage: Optional[PersistentMemory] = None
|
||||
):
|
||||
self._sessions: OrderedDict[str, MemoryEntry[SessionData]] = OrderedDict()
|
||||
self._max_sessions = max_sessions
|
||||
self._default_ttl = default_ttl
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# 持久化后端
|
||||
self._persistent = persistent_storage
|
||||
if self._persistent:
|
||||
print(f"[SessionMemory] Persistent storage enabled: {self._persistent.storage_dir}")
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
async with self._lock:
|
||||
if len(self._sessions) >= self._max_sessions and key not in self._sessions:
|
||||
self._sessions.popitem(last=False)
|
||||
|
||||
if isinstance(value, SessionData):
|
||||
session_data = value
|
||||
else:
|
||||
session_data = SessionData(session_id=key, state=value)
|
||||
|
||||
entry = MemoryEntry(
|
||||
key=key,
|
||||
value=session_data,
|
||||
ttl=ttl or self._default_ttl,
|
||||
)
|
||||
self._sessions[key] = entry
|
||||
self._sessions.move_to_end(key)
|
||||
|
||||
# 持久化存储
|
||||
if self._persistent:
|
||||
await self._persistent.store(key, session_data, ttl)
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
async with self._lock:
|
||||
# 先查内存
|
||||
entry = self._sessions.get(key)
|
||||
|
||||
# 内存未命中，查持久化
|
||||
if entry is None and self._persistent:
|
||||
persistent_data = await self._persistent.retrieve(key)
|
||||
if persistent_data:
|
||||
self._sessions[key] = MemoryEntry(
|
||||
key=key,
|
||||
value=persistent_data,
|
||||
ttl=self._default_ttl
|
||||
)
|
||||
entry = self._sessions[key]
|
||||
|
||||
if entry is None:
|
||||
return None
|
||||
|
||||
if entry.ttl and datetime.now() > entry.created_at + timedelta(seconds=entry.ttl):
|
||||
del self._sessions[key]
|
||||
if self._persistent:
|
||||
await self._persistent.delete(key)
|
||||
return None
|
||||
|
||||
self._sessions.move_to_end(key)
|
||||
return entry.value
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
async with self._lock:
|
||||
deleted = False
|
||||
if key in self._sessions:
|
||||
del self._sessions[key]
|
||||
deleted = True
|
||||
if self._persistent:
|
||||
persistent_deleted = await self._persistent.delete(key)
|
||||
deleted = deleted or persistent_deleted
|
||||
return deleted
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
async with self._lock:
|
||||
# 先查内存
|
||||
entry = self._sessions.get(key)
|
||||
if entry:
|
||||
if entry.ttl and datetime.now() > entry.created_at + timedelta(seconds=entry.ttl):
|
||||
del self._sessions[key]
|
||||
if self._persistent:
|
||||
await self._persistent.delete(key)
|
||||
return False
|
||||
return True
|
||||
|
||||
# 查持久化
|
||||
if self._persistent:
|
||||
return await self._persistent.exists(key)
|
||||
return False
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._sessions.clear()
|
||||
if self._persistent:
|
||||
await self._persistent.clear()
|
||||
|
||||
async def add_message(self, session_id: str, message: dict) -> None:
|
||||
session = await self.retrieve(session_id)
|
||||
if session is None:
|
||||
session = SessionData(session_id=session_id)
|
||||
|
||||
session.messages.append(message)
|
||||
session.updated_at = datetime.now()
|
||||
await self.store(session_id, session)
|
||||
|
||||
async def get_messages(self, session_id: str) -> list[dict]:
|
||||
session = await self.retrieve(session_id)
|
||||
return session.messages if session else []
|
||||
|
||||
async def update_state(self, session_id: str, state: dict[str, Any]) -> None:
|
||||
session = await self.retrieve(session_id)
|
||||
if session is None:
|
||||
session = SessionData(session_id=session_id)
|
||||
|
||||
session.state.update(state)
|
||||
session.updated_at = datetime.now()
|
||||
await self.store(session_id, session)
|
||||
Reference in New Issue
Block a user