90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
"""Commit normalized SOC memory items to OpenViking as structured resources."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from memory_gateway.openviking_client import OpenVikingClient
|
|
|
|
|
|
def build_resource_uri(item: dict[str, Any]) -> str:
|
|
memory_type = item.get("memory_type")
|
|
item_id = item["id"]
|
|
|
|
if memory_type == "case":
|
|
scenario = item.get("scenario", "general")
|
|
return f"viking://resources/soc-memory-poc/case/{scenario}/{item_id}.json"
|
|
|
|
if memory_type == "knowledge":
|
|
doc_type = item.get("doc_type", "general")
|
|
return f"viking://resources/soc-memory-poc/knowledge/{doc_type}/{item_id}.json"
|
|
|
|
raise ValueError(f"Unsupported memory_type for commit: {memory_type}")
|
|
|
|
|
|
def load_item(path: str | Path) -> dict[str, Any]:
|
|
path = Path(path)
|
|
with path.open("r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
async def commit_file(path: str | Path, client: OpenVikingClient) -> dict[str, Any]:
|
|
item = load_item(path)
|
|
uri = build_resource_uri(item)
|
|
result = await client.add_resource(
|
|
uri=uri,
|
|
content=json.dumps(item, ensure_ascii=False, indent=2),
|
|
resource_type="json",
|
|
wait=False,
|
|
)
|
|
return {
|
|
"path": str(path),
|
|
"uri": uri,
|
|
"result": result,
|
|
}
|
|
|
|
|
|
async def commit_directory(directory: str | Path, client: OpenVikingClient, limit: int | None = None) -> list[dict[str, Any]]:
|
|
directory = Path(directory)
|
|
paths = sorted(directory.rglob("*.json"))
|
|
if limit is not None:
|
|
paths = paths[:limit]
|
|
|
|
results: list[dict[str, Any]] = []
|
|
for path in paths:
|
|
results.append(await commit_file(path, client))
|
|
return results
|
|
|
|
|
|
async def main_async(args: argparse.Namespace) -> None:
|
|
client = OpenVikingClient()
|
|
try:
|
|
if args.path:
|
|
result = await commit_file(args.path, client)
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
else:
|
|
results = await commit_directory(args.directory, client, limit=args.limit)
|
|
print(json.dumps(results, ensure_ascii=False, indent=2))
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Commit normalized SOC items to OpenViking.")
|
|
parser.add_argument("--path", help="Single normalized JSON file to commit")
|
|
parser.add_argument("--directory", help="Directory of normalized JSON files to commit")
|
|
parser.add_argument("--limit", type=int, default=None, help="Optional limit for directory commits")
|
|
args = parser.parse_args()
|
|
|
|
if not args.path and not args.directory:
|
|
parser.error("Either --path or --directory is required")
|
|
|
|
asyncio.run(main_async(args))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|