87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
|
import pathlib
|
||
|
|
|
||
|
|
import click
|
||
|
|
from rich.console import Console
|
||
|
|
from rich.table import Table
|
||
|
|
|
||
|
|
from ontocast.onto.state import AgentState
|
||
|
|
|
||
|
|
console = Console()
|
||
|
|
|
||
|
|
|
||
|
|
def get_state_files(
|
||
|
|
directory: pathlib.Path, pattern: str = "agent_state.onto.update*.json"
|
||
|
|
) -> list[pathlib.Path]:
|
||
|
|
"""Get all AgentState files matching the pattern in the directory."""
|
||
|
|
return sorted(directory.glob(pattern))
|
||
|
|
|
||
|
|
|
||
|
|
def compare_states(states: list[tuple[pathlib.Path, AgentState]]) -> None:
|
||
|
|
"""Compare states and print a table with graph lengths."""
|
||
|
|
table = Table(title="AgentState Comparison")
|
||
|
|
table.add_column("File", style="orange1")
|
||
|
|
table.add_column("Graph Facts", justify="right")
|
||
|
|
table.add_column("Current Ontology", justify="right")
|
||
|
|
table.add_column("Ontology Addendum", justify="right")
|
||
|
|
table.add_column("Success Score", justify="right")
|
||
|
|
|
||
|
|
# Sort rows by the last number in the filename
|
||
|
|
sorted_rows = sorted(
|
||
|
|
[(fp, state) for fp, state in states],
|
||
|
|
key=lambda x: int(x[0].stem.split(".")[-1])
|
||
|
|
if x[0].stem.split(".")[-1].isdigit()
|
||
|
|
else 0,
|
||
|
|
)
|
||
|
|
|
||
|
|
for fp, state in sorted_rows:
|
||
|
|
table.add_row(
|
||
|
|
str(fp.stem),
|
||
|
|
str(len(state.current_content_unit.graph)),
|
||
|
|
str(len(state.current_ontology.graph))
|
||
|
|
if state.current_ontology is not None
|
||
|
|
else "",
|
||
|
|
str(len(state.ontology_addendum.graph)),
|
||
|
|
str(state.success_score),
|
||
|
|
)
|
||
|
|
console.print(table)
|
||
|
|
|
||
|
|
|
||
|
|
@click.command()
|
||
|
|
@click.argument(
|
||
|
|
"directory",
|
||
|
|
type=click.Path(
|
||
|
|
exists=True, file_okay=False, dir_okay=True, path_type=pathlib.Path
|
||
|
|
),
|
||
|
|
)
|
||
|
|
@click.option(
|
||
|
|
"--pattern",
|
||
|
|
default="agent_state.onto.update*.json",
|
||
|
|
help="Pattern to match state files",
|
||
|
|
)
|
||
|
|
def main(directory: pathlib.Path, pattern: str):
|
||
|
|
"""Compare AgentState files in a directory."""
|
||
|
|
state_files = get_state_files(directory, pattern)
|
||
|
|
if not state_files:
|
||
|
|
console.print(
|
||
|
|
f"[red]No state files found matching "
|
||
|
|
f"pattern '{pattern}' in {directory}[/red]"
|
||
|
|
)
|
||
|
|
return
|
||
|
|
|
||
|
|
states = []
|
||
|
|
for file_path in sorted(state_files):
|
||
|
|
try:
|
||
|
|
state = AgentState.load(file_path)
|
||
|
|
states.append((file_path, state))
|
||
|
|
except Exception as e:
|
||
|
|
console.print(f"[red]Error loading {file_path}: {str(e)}[/red]")
|
||
|
|
|
||
|
|
if states:
|
||
|
|
compare_states(states)
|
||
|
|
else:
|
||
|
|
console.print("[red]No valid state files found[/red]")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|