86 lines
3.0 KiB
Python
Executable File
86 lines
3.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import urllib.error
|
|
import urllib.request
|
|
from typing import Any
|
|
|
|
DEFAULT_GATEWAY_URL = os.environ.get("SOC_MEMORY_GATEWAY_URL", "http://127.0.0.1:1934")
|
|
DEFAULT_GATEWAY_API_KEY = os.environ.get("SOC_MEMORY_GATEWAY_API_KEY", "")
|
|
|
|
URI_PREFIXES = {
|
|
"case": "viking://resources/soc-memory-poc/case",
|
|
"knowledge": "viking://resources/soc-memory-poc/knowledge",
|
|
"all": "viking://resources/soc-memory-poc",
|
|
}
|
|
|
|
|
|
def post_json(url: str, payload: dict[str, Any], api_key: str = "") -> dict[str, Any]:
|
|
data = json.dumps(payload).encode("utf-8")
|
|
req = urllib.request.Request(url, data=data, method="POST")
|
|
req.add_header("Content-Type", "application/json")
|
|
if api_key:
|
|
req.add_header("X-API-Key", api_key)
|
|
with urllib.request.urlopen(req, timeout=30) as resp:
|
|
return json.loads(resp.read().decode("utf-8"))
|
|
|
|
|
|
def canonicalize_uri(uri: str) -> str:
|
|
if ".json/" in uri:
|
|
return uri.split(".json/", 1)[0] + ".json"
|
|
return uri
|
|
|
|
|
|
def filter_results(results: list[dict[str, Any]], prefix: str) -> list[dict[str, Any]]:
|
|
deduped: dict[str, dict[str, Any]] = {}
|
|
for item in results:
|
|
uri = item.get("uri") or ""
|
|
canonical = canonicalize_uri(uri)
|
|
if not canonical.startswith(prefix):
|
|
continue
|
|
score = item.get("score") or 0
|
|
payload = dict(item)
|
|
payload["uri"] = canonical
|
|
if canonical not in deduped or score > (deduped[canonical].get("score") or 0):
|
|
deduped[canonical] = payload
|
|
return sorted(deduped.values(), key=lambda entry: entry.get("score") or 0, reverse=True)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Search SOC Memory Gateway for case / knowledge context.")
|
|
parser.add_argument("--query", required=True, help="Search query")
|
|
parser.add_argument("--kind", choices=["case", "knowledge", "all"], default="all", help="SOC resource scope")
|
|
parser.add_argument("--limit", type=int, default=5, help="Max results")
|
|
parser.add_argument("--gateway-url", default=DEFAULT_GATEWAY_URL, help="Memory Gateway base URL")
|
|
parser.add_argument("--api-key", default=DEFAULT_GATEWAY_API_KEY, help="Gateway API key if required")
|
|
args = parser.parse_args()
|
|
|
|
prefix = URI_PREFIXES[args.kind]
|
|
payload = {
|
|
"query": args.query,
|
|
"limit": max(args.limit * 5, 10),
|
|
"uri": prefix,
|
|
}
|
|
try:
|
|
result = post_json(args.gateway_url.rstrip("/") + "/api/search", payload, api_key=args.api_key)
|
|
except urllib.error.URLError as exc:
|
|
raise SystemExit(f"Gateway search failed: {exc}") from exc
|
|
|
|
raw_results = result.get("results", [])
|
|
filtered = filter_results(raw_results, prefix)
|
|
output = {
|
|
"query": args.query,
|
|
"kind": args.kind,
|
|
"uri_prefix": prefix,
|
|
"results": filtered[: args.limit],
|
|
"total": len(filtered),
|
|
}
|
|
print(json.dumps(output, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|