import asyncio
import json
import logging
import time
from pathlib import Path
from typing import List, Dict, Any, Optional
from everos import AsyncEverOS, NotFoundError
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
async def wait_for_task(
client: AsyncEverOS,
task_id: str,
poll_interval: float = 2.0,
max_attempts: int = 30,
) -> str:
"""Poll a task until completion. Returns final status."""
for attempt in range(max_attempts):
try:
resp = await client.v1.tasks.retrieve(task_id)
except NotFoundError:
# Task completed and was cleaned up (short TTL)
logger.info(f" [attempt {attempt + 1}] task {task_id} done (expired)")
return "done"
status = resp.data.status if resp.data else "unknown"
logger.info(f" [attempt {attempt + 1}] task {task_id} status={status}")
if status in ("success", "failed", "completed", "done"):
return status
await asyncio.sleep(poll_interval)
return "timeout"
class BatchImporter:
"""Batch import conversations into EverOS using the v1 SDK."""
def __init__(
self,
concurrency: int = 10,
chunk_size: int = 50,
):
self.client = AsyncEverOS()
self.concurrency = concurrency
self.chunk_size = chunk_size
self.semaphore = asyncio.Semaphore(concurrency)
# Statistics
self.stats = {
"messages_sent": 0,
"messages_failed": 0,
"conversations_processed": 0,
}
async def import_file(self, file_path: str) -> Dict[str, Any]:
"""Import a single conversation JSON file."""
with open(file_path, "r") as f:
data = json.load(f)
return await self.import_conversation(data)
async def import_conversation(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""Import a conversation (group or personal)."""
is_group = "group_id" in data
conversation_id = data.get("group_id") or data.get("user_id")
messages = data.get("messages", [])
logger.info(
f"Processing {'group' if is_group else 'personal'} "
f"conversation: {conversation_id} ({len(messages)} messages)"
)
# 1. Set up group metadata if needed
if is_group and data.get("group_meta"):
await self._setup_group(data["group_id"], data["group_meta"])
# 2. Import messages in chunks
total_sent = 0
total_failed = 0
for i in range(0, len(messages), self.chunk_size):
chunk = messages[i : i + self.chunk_size]
sent, failed = await self._send_chunk(data, chunk, is_group)
total_sent += sent
total_failed += failed
# 3. Flush to trigger memory extraction
await self._flush(data, is_group)
self.stats["messages_sent"] += total_sent
self.stats["messages_failed"] += total_failed
self.stats["conversations_processed"] += 1
logger.info(
f"Completed {conversation_id}: {total_sent} sent, {total_failed} failed"
)
return {
"conversation_id": conversation_id,
"messages_sent": total_sent,
"messages_failed": total_failed,
}
async def _setup_group(self, group_id: str, group_meta: Dict):
"""Create or update group metadata."""
try:
await self.client.v1.groups.create(
group_id=group_id,
name=group_meta.get("name", group_id),
description=group_meta.get("description", ""),
)
except Exception as e:
logger.warning(f"Failed to create group {group_id}: {e}")
async def _send_chunk(
self, data: Dict, chunk: List[Dict], is_group: bool
) -> tuple:
"""Send a chunk of messages using async mode."""
async with self.semaphore:
try:
if is_group:
response = await self.client.v1.memories.group.add(
group_id=data["group_id"],
async_mode=True,
messages=chunk,
)
else:
response = await self.client.v1.memories.add(
user_id=data["user_id"],
async_mode=True,
messages=chunk,
)
# Wait for async task to complete
task_id = response.data.task_id if response.data else None
if task_id:
status = await wait_for_task(self.client, task_id)
if status in ("failed", "timeout"):
logger.error(f"Task {task_id} ended with status: {status}")
return 0, len(chunk)
return len(chunk), 0
except Exception as e:
logger.error(f"Failed to send chunk: {e}")
return 0, len(chunk)
async def _flush(self, data: Dict, is_group: bool):
"""Flush to trigger memory extraction."""
try:
if is_group:
await self.client.v1.memories.group.flush(
group_id=data["group_id"]
)
else:
await self.client.v1.memories.flush(user_id=data["user_id"])
except Exception as e:
logger.warning(f"Flush failed: {e}")
async def import_directory(self, directory: str) -> List[Dict]:
"""Import all JSON files in a directory."""
path = Path(directory)
files = sorted(path.glob("*.json"))
logger.info(f"Found {len(files)} files to import")
results = []
for file_path in files:
try:
result = await self.import_file(str(file_path))
results.append(result)
except Exception as e:
logger.error(f"Failed to import {file_path}: {e}")
results.append({"file": str(file_path), "error": str(e)})
return results
def get_stats(self) -> Dict[str, int]:
"""Get import statistics."""
return self.stats.copy()
# Usage
async def main():
importer = BatchImporter(
concurrency=20,
chunk_size=50,
)
# Import a single file
result = await importer.import_file("conversation_001.json")
print(f"Result: {result}")
# Import a directory
results = await importer.import_directory("./exports/")
print(f"Imported {len(results)} conversations")
print(f"Stats: {importer.get_stats()}")
if __name__ == "__main__":
asyncio.run(main())
