From 39a1bda16b074ac59c9f9c7ff09d4f47988d067e Mon Sep 17 00:00:00 2001 From: Cameron Pfiffer Date: Wed, 14 May 2025 22:10:36 -0700 Subject: [PATCH] feat: Implement KuzuDB integration for ATProto records and enhance database schema --- .gitignore | 5 +- KuzuDB_README.md | 124 +++++++ requirements.txt | 2 + src/db_manager.py | 737 +++++++++++++++++++++++++++++++++++++- src/jetstream_consumer.py | 78 +++- test_db.py | 207 +++++++++++ test_kuzu.py | 165 +++++++++ 7 files changed, 1308 insertions(+), 10 deletions(-) create mode 100644 KuzuDB_README.md create mode 100644 test_db.py create mode 100644 test_kuzu.py diff --git a/.gitignore b/.gitignore index a6646f7..c4a752e 100644 --- a/.gitignore +++ b/.gitignore @@ -7,9 +7,6 @@ user_info_cache.json *.pyc README_files/ README.html -<<<<<<< HEAD **/CLAUDE.local.md -======= -.aider* ->>>>>>> 391d88205f72434bf84f317e9452998bf10cebc7 +demo_db/ diff --git a/KuzuDB_README.md b/KuzuDB_README.md new file mode 100644 index 0000000..77e69ba --- /dev/null +++ b/KuzuDB_README.md @@ -0,0 +1,124 @@ +# KuzuDB Integration for Comind + +This document describes the KuzuDB integration for tracking ATProto records in Comind. + +## Overview + +KuzuDB is an embedded graph database that is used to store and query ATProto records for Comind. The integration includes: + +- Storing ATRecord nodes with metadata and content +- Tracking relationships between records +- Applying appropriate labels to different record types +- Full-text search capabilities +- Query functionality to explore the graph + +## Schema Design + +The database schema includes the following node tables: + +- **Repo**: Represents a repository owned by a user, identified by DID +- **ATRecord**: All ATProto records with their metadata and content +- **User**: User information (existing) +- **Record**: Base record type (existing) +- **Sphere**: Sphere record (existing) +- **BlipConcept**, **BlipEmotion**, **BlipThought**: Specialized record types (existing) + +And relationship tables: + +- **OWNS**: Connects a Repo to its ATRecords +- **FOLLOWS**: Connects repos when users follow each other +- **LIKES**: Connects an ATRecord to another when a like occurs +- **REPOSTS**: Connects an ATRecord to another when a repost occurs +- **BLOCKS**: Connects repos when a user blocks another +- **LINKS**: General relationship between records +- **IN_SPHERE**: Connects records to spheres +- **AUTHORED**: Connects users to records (existing) +- **TARGET**: Connects records to their targets (existing) + +## Usage + +### Starting the Jetstream Consumer with KuzuDB + +The jetstream consumer can now store all ATProto records in KuzuDB as they arrive. Use the following command to run it: + +```bash +python src/jetstream_consumer.py --db-path ./your_db_path --comind your_comind_name +``` + +Options: +- `--db-path`: Path to the KuzuDB database directory (default: ./demo_db) +- `--disable-db`: Disable storing ATProto records in KuzuDB + +### Directly Using the Database Manager + +You can also use the DBManager directly in your code: + +```python +from src.db_manager import DBManager + +# Initialize the manager +db = DBManager('./your_db_path') + +# Store an ATProto record +db.store_atproto_record( + uri="at://did:example/app.bsky.feed.post/abcde", + cid="bafyreihgxxx", + nsid="app.bsky.feed.post", + record={"text": "Hello world", "createdAt": "2023-01-01T00:00:00Z"}, + author_did="did:example", + rkey="abcde", + labels=["ATRecord", "Post"] +) + +# Query records +posts = db.list_atproto_records(nsid="app.bsky.feed.post", limit=10) + +# Find relationships +relationships = db.query_atproto_relationships(source_uri="at://did:example/app.bsky.feed.post/abcde") +``` + +### Testing the Integration + +Run the included test script to verify the database is working: + +```bash +python test_kuzu.py +``` + +This will show all stored records, count them by type, display relationships, and test the search functionality. + +## Node Labels + +Records in the database are automatically assigned labels based on their type: + +- All records: `ATRecord` +- Posts: `ATRecord`, `Post` +- Concepts: `ATRecord`, `Blip`, `Concept` +- Thoughts: `ATRecord`, `Blip`, `Thought` +- Emotions: `ATRecord`, `Blip`, `Emotion` +- Spheres: `ATRecord`, `Core` + +## Querying the Database + +The KuzuDB integration provides several methods for querying: + +- `list_atproto_records()`: List records with optional filtering +- `get_atproto_record()`: Get a single record by URI +- `query_atproto_relationships()`: Find relationships between records +- `find_similar_records()`: Search records by text content + +Internally, the database executes Cypher queries, which can be extended as needed. + +## Extending the Integration + +To add support for new record types: + +1. Update the record processing in `store_atproto_record()` to handle the new type +2. Add appropriate label assignment +3. Add any specialized relationships + +## Troubleshooting + +- If you see errors about "table already exists" when setting up the schema, they can be safely ignored +- If full-text search doesn't work, the basic search fallback will be used automatically +- Make sure the database directory exists and is writable \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 118c7aa..fa3dc82 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,12 +13,14 @@ httpcore==1.0.7 httpx==0.28.1 idna==3.10 jiter==0.9.0 +kuzu==0.0.6 libipld==3.0.1 linkify-it-py==2.0.3 markdown-it-py==3.0.0 mdit-py-plugins==0.4.2 mdurl==0.1.2 openai==1.69.0 +pandas>=2.0.0 platformdirs==4.3.7 pycparser==2.22 pydantic==2.11.1 diff --git a/src/db_manager.py b/src/db_manager.py index 3b89b7d..415a7f1 100644 --- a/src/db_manager.py +++ b/src/db_manager.py @@ -69,6 +69,30 @@ class DBManager: ) """) + # Create Repo node table + self.conn.execute(""" + CREATE NODE TABLE Repo ( + did STRING PRIMARY KEY, + handle STRING, + receivedAt TIMESTAMP + ) + """) + + # Create ATRecord node table (base for all record types) + self.conn.execute(""" + CREATE NODE TABLE ATRecord ( + uri STRING PRIMARY KEY, + cid STRING, + nsid STRING, + rkey STRING, + createdAt TIMESTAMP, + receivedAt TIMESTAMP, + raw STRING, + text STRING, + labels STRING + ) + """) + # Create Record node table (base for all record types) self.conn.execute(""" CREATE NODE TABLE Record ( @@ -133,6 +157,14 @@ class DBManager: ) """) + # OWNS relationship between Repo and ATRecord + self.conn.execute(""" + CREATE REL TABLE OWNS ( + FROM Repo TO ATRecord, + createdAt TIMESTAMP + ) + """) + # IN_SPHERE relationship between Record and Sphere self.conn.execute(""" CREATE REL TABLE IN_SPHERE ( @@ -160,6 +192,42 @@ class DBManager: ) """) + # FOLLOWS relationship between ATRecords + self.conn.execute(""" + CREATE REL TABLE FOLLOWS ( + FROM ATRecord TO ATRecord, + sourceRecord STRING, + createdAt TIMESTAMP + ) + """) + + # LIKES relationship between ATRecords + self.conn.execute(""" + CREATE REL TABLE LIKES ( + FROM ATRecord TO ATRecord, + sourceRecord STRING, + createdAt TIMESTAMP + ) + """) + + # REPOSTS relationship between ATRecords + self.conn.execute(""" + CREATE REL TABLE REPOSTS ( + FROM ATRecord TO ATRecord, + sourceRecord STRING, + createdAt TIMESTAMP + ) + """) + + # BLOCKS relationship between ATRecords + self.conn.execute(""" + CREATE REL TABLE BLOCKS ( + FROM ATRecord TO ATRecord, + sourceRecord STRING, + createdAt TIMESTAMP + ) + """) + logger.info("Successfully created database schema") except Exception as e: @@ -901,6 +969,591 @@ class DBManager: except Exception as e: logger.error(f"Error in basic record search: {str(e)}") return [] + + def store_atproto_record(self, uri: str, cid: str, nsid: str, record: Dict, + author_did: str, rkey: str, labels: List[str] = None): + """ + Store an ATProto record from the jetstream in the database. + + Args: + uri: The URI of the record + cid: The content identifier of the record + nsid: The NSID (namespace ID) of the record (e.g., app.bsky.feed.post) + record: The record data as a dictionary + author_did: The DID of the record's author/repo owner + rkey: The record key identifier + labels: Additional labels to apply to the node (e.g., 'Blip', 'Concept') + """ + try: + # Convert record to JSON string + record_json = json.dumps(record) + + # Extract text content if available (depends on record type) + text = "" + if nsid == "app.bsky.feed.post" and "text" in record: + text = record.get("text", "") + elif "me.comind" in nsid: + # Extract text from different comind record types + if "blip.concept" in nsid and "generated" in record and "text" in record["generated"]: + text = record["generated"]["text"] + elif "blip.thought" in nsid and "generated" in record and "text" in record["generated"]: + text = record["generated"]["text"] + elif "blip.emotion" in nsid and "generated" in record and "text" in record["generated"]: + text = record["generated"]["text"] + elif "sphere.core" in nsid: + text = record.get("text", "") + + # Format dates + created_at = datetime.now() + if "createdAt" in record: + try: + if isinstance(record["createdAt"], str): + created_at = datetime.fromisoformat(record["createdAt"].replace('Z', '+00:00')) + except Exception as e: + logger.warning(f"Error parsing createdAt timestamp: {e}. Using current time.") + + received_at = datetime.now() + + # Convert labels to string for storage + labels_str = "" + if labels: + labels_str = ",".join(labels) + + # First, ensure the repo exists + self.conn.execute(""" + MERGE (r:Repo {did: $did}) + SET r.receivedAt = $receivedAt + """, { + 'did': author_did, + 'receivedAt': received_at + }) + + # Store the ATRecord + self.conn.execute(""" + MERGE (r:ATRecord {uri: $uri}) + SET r.cid = $cid, + r.nsid = $nsid, + r.rkey = $rkey, + r.createdAt = $createdAt, + r.receivedAt = $receivedAt, + r.raw = $raw, + r.text = $text, + r.labels = $labels + """, { + 'uri': uri, + 'cid': cid, + 'nsid': nsid, + 'rkey': rkey, + 'createdAt': created_at, + 'receivedAt': received_at, + 'raw': record_json, + 'text': text, + 'labels': labels_str + }) + + # Create OWNS relationship between repo and record + self.conn.execute(""" + MATCH (repo:Repo {did: $did}) + MATCH (record:ATRecord {uri: $uri}) + MERGE (repo)-[rel:OWNS]->(record) + SET rel.createdAt = $createdAt + """, { + 'did': author_did, + 'uri': uri, + 'createdAt': received_at + }) + + # Handle special record types that create relationships + if nsid == "app.bsky.graph.follow" and "subject" in record: + target_did = record["subject"] + # Find the target repo + self.conn.execute(""" + MERGE (target:Repo {did: $target_did}) + """, { + 'target_did': target_did + }) + + # Create FOLLOWS relationship + self.conn.execute(""" + MATCH (source:ATRecord {uri: $uri}) + MATCH (source_repo:Repo {did: $source_did}) + MATCH (target_repo:Repo {did: $target_did}) + MERGE (source_repo)-[rel:FOLLOWS]->(target_repo) + SET rel.sourceRecord = $uri, + rel.createdAt = $createdAt + """, { + 'uri': uri, + 'source_did': author_did, + 'target_did': target_did, + 'createdAt': created_at + }) + + elif nsid == "app.bsky.feed.like" and "subject" in record: + target_uri = record["subject"]["uri"] + # Create LIKES relationship + self.conn.execute(""" + MATCH (source:ATRecord {uri: $uri}) + MATCH (target:ATRecord {uri: $target_uri}) + MERGE (source)-[rel:LIKES]->(target) + SET rel.sourceRecord = $uri, + rel.createdAt = $createdAt + """, { + 'uri': uri, + 'target_uri': target_uri, + 'createdAt': created_at + }) + + elif nsid == "app.bsky.feed.repost" and "subject" in record: + target_uri = record["subject"]["uri"] + # Create REPOSTS relationship + self.conn.execute(""" + MATCH (source:ATRecord {uri: $uri}) + MATCH (target:ATRecord {uri: $target_uri}) + MERGE (source)-[rel:REPOSTS]->(target) + SET rel.sourceRecord = $uri, + rel.createdAt = $createdAt + """, { + 'uri': uri, + 'target_uri': target_uri, + 'createdAt': created_at + }) + + elif nsid == "app.bsky.graph.block" and "subject" in record: + target_did = record["subject"] + # Find the target repo + self.conn.execute(""" + MERGE (target:Repo {did: $target_did}) + """, { + 'target_did': target_did + }) + + # Create BLOCKS relationship + self.conn.execute(""" + MATCH (source:ATRecord {uri: $uri}) + MATCH (source_repo:Repo {did: $source_did}) + MATCH (target_repo:Repo {did: $target_did}) + MERGE (source_repo)-[rel:BLOCKS]->(target_repo) + SET rel.sourceRecord = $uri, + rel.createdAt = $createdAt + """, { + 'uri': uri, + 'source_did': author_did, + 'target_did': target_did, + 'createdAt': created_at + }) + + # Handle Comind specific records + if "me.comind" in nsid: + # Add appropriate label based on record type + if "blip.concept" in nsid: + self.conn.execute(""" + MATCH (r:ATRecord {uri: $uri}) + SET r.labels = CASE + WHEN r.labels = '' THEN 'ATRecord,Blip,Concept' + WHEN r.labels CONTAINS 'Concept' THEN r.labels + ELSE r.labels + ',Blip,Concept' + END + """, {'uri': uri}) + + elif "blip.thought" in nsid: + self.conn.execute(""" + MATCH (r:ATRecord {uri: $uri}) + SET r.labels = CASE + WHEN r.labels = '' THEN 'ATRecord,Blip,Thought' + WHEN r.labels CONTAINS 'Thought' THEN r.labels + ELSE r.labels + ',Blip,Thought' + END + """, {'uri': uri}) + + elif "blip.emotion" in nsid: + self.conn.execute(""" + MATCH (r:ATRecord {uri: $uri}) + SET r.labels = CASE + WHEN r.labels = '' THEN 'ATRecord,Blip,Emotion' + WHEN r.labels CONTAINS 'Emotion' THEN r.labels + ELSE r.labels + ',Blip,Emotion' + END + """, {'uri': uri}) + + elif "sphere.core" in nsid: + self.conn.execute(""" + MATCH (r:ATRecord {uri: $uri}) + SET r.labels = CASE + WHEN r.labels = '' THEN 'ATRecord,Core' + WHEN r.labels CONTAINS 'Core' THEN r.labels + ELSE r.labels + ',Core' + END + """, {'uri': uri}) + + # Handle "from" records that create references + if "from" in record and isinstance(record["from"], list): + for ref in record["from"]: + if "uri" in ref: + from_uri = ref["uri"] + self.conn.execute(""" + MATCH (target:ATRecord {uri: $uri}) + MERGE (source:ATRecord {uri: $from_uri}) + MERGE (source)-[rel:LINKS]->(target) + SET rel.relType = 'REFERENCES', + rel.createdAt = $createdAt + """, { + 'uri': uri, + 'from_uri': from_uri, + 'createdAt': created_at + }) + + # Handle relationship links + if "relationship.link" in nsid: + if "target" in record: + target_uri = record["target"] + self.conn.execute(""" + MATCH (source:ATRecord {uri: $uri}) + MERGE (target:ATRecord {uri: $target_uri}) + MERGE (source)-[rel:LINKS]->(target) + SET rel.relType = $rel_type, + rel.strength = $strength, + rel.note = $note, + rel.createdAt = $createdAt + """, { + 'uri': uri, + 'target_uri': target_uri, + 'rel_type': record.get("relationship", "REFERENCES"), + 'strength': record.get("strength", 1.0), + 'note': record.get("note", ""), + 'createdAt': created_at + }) + + # Handle sphere assignment + if "relationship.sphere" in nsid and "sphere_uri" in record and "target" in record: + sphere_uri = record["sphere_uri"] + target_uri = record["target"]["uri"] + + self.conn.execute(""" + MATCH (record:ATRecord {uri: $target_uri}) + MATCH (sphere:ATRecord {uri: $sphere_uri}) + MERGE (sphere)-[rel:CONGREGATES]->(record) + SET rel.createdAt = $createdAt + """, { + 'target_uri': target_uri, + 'sphere_uri': sphere_uri, + 'createdAt': created_at + }) + + logger.debug(f"Stored ATProto record: {nsid}/{rkey}") + return True + + except Exception as e: + logger.error(f"Error storing ATProto record {uri}: {str(e)}") + logger.exception(e) + return False + + def get_atproto_record(self, uri: str) -> Optional[Dict]: + """ + Retrieve an ATProto record by its URI. + + Args: + uri: The URI of the record to retrieve + + Returns: + The record as a dictionary if found, None otherwise + """ + try: + query = """ + MATCH (r:ATRecord {uri: $uri}) + RETURN r.uri as uri, r.cid as cid, r.nsid as nsid, r.raw as raw, + r.createdAt as createdAt, r.receivedAt as receivedAt, r.labels as labels + """ + + result = self.conn.execute(query, {'uri': uri}) + + if result.has_next(): + row = result.get_next() + uri, cid, nsid, raw, created_at, received_at, labels = row + + record = { + 'uri': uri, + 'cid': cid, + 'nsid': nsid, + 'value': json.loads(raw), + 'createdAt': created_at, + 'receivedAt': received_at, + 'labels': labels.split(',') if labels else [] + } + return record + else: + return None + + except Exception as e: + logger.error(f"Error retrieving ATProto record {uri}: {str(e)}") + return None + + def list_atproto_records(self, nsid: str = None, labels: List[str] = None, limit: int = 100) -> List[Dict]: + """ + List ATProto records with optional filtering by nsid and labels. + + Args: + nsid: Optional namespace ID to filter by + labels: Optional list of labels to filter by + limit: Maximum number of records to return + + Returns: + List of matching records + """ + try: + where_clauses = [] + params = {'limit': limit} + + if nsid: + where_clauses.append("r.nsid = $nsid") + params['nsid'] = nsid + + if labels and len(labels) > 0: + label_conditions = [] + for i, label in enumerate(labels): + param_name = f'label_{i}' + label_conditions.append(f"r.labels CONTAINS ${param_name}") + params[param_name] = label + + where_clauses.append(f"({' AND '.join(label_conditions)})") + + where_clause = " AND ".join(where_clauses) if where_clauses else "1=1" + + query = f""" + MATCH (r:ATRecord) + WHERE {where_clause} + RETURN r.uri as uri, r.cid as cid, r.nsid as nsid, r.raw as raw, + r.createdAt as createdAt, r.receivedAt as receivedAt, r.labels as labels + ORDER BY r.createdAt DESC + LIMIT $limit + """ + + result = self.conn.execute(query, params) + records = [] + + while result.has_next(): + row = result.get_next() + uri, cid, nsid, raw, created_at, received_at, labels_str = row + + record = { + 'uri': uri, + 'cid': cid, + 'nsid': nsid, + 'value': json.loads(raw), + 'createdAt': created_at, + 'receivedAt': received_at, + 'labels': labels_str.split(',') if labels_str else [] + } + records.append(record) + + return records + + except Exception as e: + logger.error(f"Error listing ATProto records: {str(e)}") + return [] + + def query_atproto_relationships(self, source_uri: str, relationship_type: str = None, limit: int = 100) -> List[Dict]: + """ + Query relationships from a source ATProto record. + + Args: + source_uri: The URI of the source record + relationship_type: Optional relationship type to filter by (FOLLOWS, LIKES, REPOSTS, etc.) + limit: Maximum number of relationships to return + + Returns: + List of related records with relationship information + """ + try: + relationships = [] + params = {'source_uri': source_uri, 'limit': limit} + + # Determine the relationship table based on type + if relationship_type == "FOLLOWS": + query = """ + MATCH (source:ATRecord {uri: $source_uri}) + MATCH (source_repo:Repo)<-[:OWNS]-(source) + MATCH (source_repo)-[rel:FOLLOWS]->(target_repo:Repo) + MATCH (target_repo)-[:OWNS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'FOLLOWS' as rel_type + LIMIT $limit + """ + elif relationship_type == "LIKES": + query = """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:LIKES]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'LIKES' as rel_type + LIMIT $limit + """ + elif relationship_type == "REPOSTS": + query = """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:REPOSTS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'REPOSTS' as rel_type + LIMIT $limit + """ + elif relationship_type == "BLOCKS": + query = """ + MATCH (source:ATRecord {uri: $source_uri}) + MATCH (source_repo:Repo)<-[:OWNS]-(source) + MATCH (source_repo)-[rel:BLOCKS]->(target_repo:Repo) + MATCH (target_repo)-[:OWNS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'BLOCKS' as rel_type + LIMIT $limit + """ + elif relationship_type == "LINKS": + query = """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:LINKS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + '' as record_uri, rel.createdAt as created_at, + rel.relType as rel_type, rel.strength as strength, rel.note as note + LIMIT $limit + """ + elif relationship_type == "IN_SPHERE": + query = """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:IN_SPHERE]->(sphere:ATRecord) + RETURN source.uri as source_uri, sphere.uri as target_uri, + '' as record_uri, rel.createdAt as created_at, + 'IN_SPHERE' as rel_type + LIMIT $limit + """ + else: + # Query all relationship types + queries = [ + """ + MATCH (source:ATRecord {uri: $source_uri}) + MATCH (source_repo:Repo)<-[:OWNS]-(source) + MATCH (source_repo)-[rel:FOLLOWS]->(target_repo:Repo) + MATCH (target_repo)-[:OWNS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'FOLLOWS' as rel_type + LIMIT $limit + """, + """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:LIKES]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'LIKES' as rel_type + LIMIT $limit + """, + """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:REPOSTS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'REPOSTS' as rel_type + LIMIT $limit + """, + """ + MATCH (source:ATRecord {uri: $source_uri}) + MATCH (source_repo:Repo)<-[:OWNS]-(source) + MATCH (source_repo)-[rel:BLOCKS]->(target_repo:Repo) + MATCH (target_repo)-[:OWNS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + rel.sourceRecord as record_uri, rel.createdAt as created_at, + 'BLOCKS' as rel_type + LIMIT $limit + """, + """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:LINKS]->(target:ATRecord) + RETURN source.uri as source_uri, target.uri as target_uri, + '' as record_uri, rel.createdAt as created_at, + rel.relType as rel_type, rel.strength as strength, rel.note as note + LIMIT $limit + """, + """ + MATCH (source:ATRecord {uri: $source_uri})-[rel:IN_SPHERE]->(sphere:ATRecord) + RETURN source.uri as source_uri, sphere.uri as target_uri, + '' as record_uri, rel.createdAt as created_at, + 'IN_SPHERE' as rel_type + LIMIT $limit + """ + ] + + # Execute each query and combine results + for q in queries: + try: + result = self.conn.execute(q, params) + while result.has_next(): + row = result.get_next() + if len(row) == 5: + source_uri, target_uri, record_uri, created_at, rel_type = row + rel = { + 'source_uri': source_uri, + 'target_uri': target_uri, + 'record_uri': record_uri, + 'created_at': created_at, + 'type': rel_type + } + else: + source_uri, target_uri, record_uri, created_at, rel_type, strength, note = row + rel = { + 'source_uri': source_uri, + 'target_uri': target_uri, + 'record_uri': record_uri, + 'created_at': created_at, + 'type': rel_type, + 'strength': strength, + 'note': note + } + relationships.append(rel) + + if len(relationships) >= limit: + break + + if len(relationships) >= limit: + break + + except Exception as e: + logger.error(f"Error executing relationship query: {str(e)}") + continue + + # Sort by created_at + relationships.sort(key=lambda r: r.get('created_at', ''), reverse=True) + + # Limit to requested number + relationships = relationships[:limit] + + return relationships + + # Execute the specific relationship query if a type was specified + if relationship_type: + result = self.conn.execute(query, params) + + while result.has_next(): + row = result.get_next() + if len(row) == 5: + source_uri, target_uri, record_uri, created_at, rel_type = row + rel = { + 'source_uri': source_uri, + 'target_uri': target_uri, + 'record_uri': record_uri, + 'created_at': created_at, + 'type': rel_type + } + else: + source_uri, target_uri, record_uri, created_at, rel_type, strength, note = row + rel = { + 'source_uri': source_uri, + 'target_uri': target_uri, + 'record_uri': record_uri, + 'created_at': created_at, + 'type': rel_type, + 'strength': strength, + 'note': note + } + relationships.append(rel) + + return relationships + + except Exception as e: + logger.error(f"Error querying ATProto relationships for {source_uri}: {str(e)}") + return [] # Function to integrate with record_manager.py @@ -945,4 +1598,86 @@ def mirror_record_to_db(record_manager, db_manager, collection: str, record: Dic sphere_uri=sphere_uri ) - return uri \ No newline at end of file + return uri + +# Helper function to process ATProto events from the jetstream +def process_atproto_event(db_manager, event: Dict) -> bool: + """ + Process an ATProto event from the jetstream and store it in the database. + + Args: + db_manager: The DBManager instance + event: The event data from the jetstream + + Returns: + True if the event was processed successfully, False otherwise + """ + try: + # Extract event data + author_did = event.get("did") + if not author_did: + logger.error("No author DID found in event") + return False + + # Handle create operation + if (event.get("kind") == "commit" and + event.get("commit", {}).get("operation") == "create"): + + # Extract record data + commit = event.get("commit", {}) + collection = commit.get("collection") + rkey = commit.get("rkey") + record = commit.get("record", {}) + cid = commit.get("cid", "") + + # Skip if missing required data + if not collection or not rkey or not record: + logger.warning(f"Missing required data in event: collection={collection}, rkey={rkey}") + return False + + # Construct URI + uri = f"at://{author_did}/{collection}/{rkey}" + + # Determine labels based on collection + labels = ["ATRecord"] + + if collection == "app.bsky.feed.post": + labels.append("Post") + elif "me.comind" in collection: + if "blip.concept" in collection: + labels.extend(["Blip", "Concept"]) + elif "blip.thought" in collection: + labels.extend(["Blip", "Thought"]) + elif "blip.emotion" in collection: + labels.extend(["Blip", "Emotion"]) + elif "sphere.core" in collection: + labels.append("Core") + + # Store the record + success = db_manager.store_atproto_record( + uri=uri, + cid=cid, + nsid=collection, + record=record, + author_did=author_did, + rkey=rkey, + labels=labels + ) + + return success + + # Handle other operations (delete, update) + elif (event.get("kind") == "commit" and + event.get("commit", {}).get("operation") in ["delete", "update"]): + + # TODO: Implement handling for delete and update operations + logger.info(f"Event operation {event.get('commit', {}).get('operation')} not yet implemented") + return False + + else: + logger.warning(f"Unknown event kind or operation: {event.get('kind')} / {event.get('commit', {}).get('operation')}") + return False + + except Exception as e: + logger.error(f"Error processing ATProto event: {str(e)}") + return False \ No newline at end of file diff --git a/src/jetstream_consumer.py b/src/jetstream_consumer.py index 0e3de1b..2d21ede 100644 --- a/src/jetstream_consumer.py +++ b/src/jetstream_consumer.py @@ -3,7 +3,7 @@ from pydantic import BaseModel, Field from datetime import datetime import json import asyncio -from typing import List, Optional, Set, Any +from typing import List, Optional, Set, Any, Dict import websockets import time import logging @@ -27,6 +27,7 @@ from src.lexicon_utils import ( ) import src.structured_gen as structured_gen from src.record_manager import RecordManager +from src.db_manager import DBManager, process_atproto_event from atproto_client import Client import yaml @@ -272,6 +273,8 @@ async def process_event( thread_depth: int = 15, user_info_cache: UserInfoCache = None, comind: Comind = None, + db_manager: DBManager = None, + original_event: Dict = None # Add parameter to receive the original event ) -> None: """Process an event and generate thoughts, emotions, and concepts for it""" try: @@ -430,6 +433,41 @@ async def process_event( from_refs=list_of_strong_refs # Pass the collected references ) + # Also store the event in KuzuDB if a db_manager is provided + if db_manager is not None: + try: + # If we have the original event object, use it + if original_event: + process_atproto_event(db_manager, original_event) + else: + # Otherwise reconstruct a minimal event from the parameters we have + # Extract rkey from post_uri (last part after the last /) + rkey = post_uri.split('/')[-1] + # Extract collection from event_kind + collection = event_kind + + # Get the record data from the target_post + record = {} + if hasattr(target_post, 'record'): + record = target_post.record.model_dump() + + # Construct an event object that process_atproto_event can handle + reconstructed_event = { + "did": author_did, + "kind": "commit", + "commit": { + "operation": "create", + "collection": collection, + "rkey": rkey, + "record": record, + "cid": post_cid + } + } + process_atproto_event(db_manager, reconstructed_event) + except Exception as db_error: + logger.error(f"Error storing event in KuzuDB: {db_error}") + # Don't raise the exception - we still want to process the event for comind + except Exception as e: logger.error(f"Error processing post {post_uri}: {e}") raise e @@ -439,7 +477,8 @@ async def connect_to_jetstream( activated_dids_file: str, jetstream_host: str = JETSTREAM_HOST, thread_depth: int = 15, - comind: Comind = None + comind: Comind = None, + db_manager: DBManager = None ) -> None: """Connect to Jetstream and process incoming messages""" global activated_dids @@ -516,6 +555,14 @@ async def connect_to_jetstream( logger.warning(f"No author DID found in event: {event}") raise Exception(f"No author DID found in event: {event}") + # Also store the event in KuzuDB if a db_manager is provided + if db_manager is not None: + try: + process_atproto_event(db_manager, event) + except Exception as db_error: + logger.error(f"Error storing event in KuzuDB: {db_error}") + # Don't raise the exception - we still want to process the event for comind + # Check if it's a post creation event if (event.get("kind") == "commit" and event.get("commit", {}).get("operation") == "create"): @@ -546,7 +593,9 @@ async def connect_to_jetstream( root_post_uri=root_post_uri, thread_depth=thread_depth, user_info_cache=user_info_cache, - comind=comind + comind=comind, + db_manager=db_manager, + original_event=event ) elif collection == "app.bsky.feed.like": # Extract post URI and CID @@ -562,7 +611,9 @@ async def connect_to_jetstream( post_cid, thread_depth=thread_depth, user_info_cache=user_info_cache, - comind=comind + comind=comind, + db_manager=db_manager, + original_event=event ) else: logger.warning(f"Unknown collection message received: {collection}") @@ -639,6 +690,10 @@ async def main(): help="Sphere to attach comind records to. Default is to not use a sphere.") parser.add_argument("--comind", "-c", type=str, default=None, help="Comind to use for processing. Required.") + parser.add_argument("--db-path", type=str, default="./demo_db", + help="Path to the Kuzu database directory. Default is './demo_db'.") + parser.add_argument("--disable-db", action="store_true", + help="Disable storing ATProto records in Kuzu database.") args = parser.parse_args() @@ -728,6 +783,18 @@ async def main(): comind.core_perspective = sphere_to_use.value["text"] else: logger.warning("No sphere provided. Comind will not be attached to any sphere.") + + # Initialize the database manager if not disabled + db_manager = None + if not args.disable_db: + try: + logger.info(f"Initializing Kuzu database at {args.db_path}") + db_manager = DBManager(args.db_path) + db_manager.setup_schema() + logger.info("Kuzu database initialized successfully") + except Exception as e: + logger.error(f"Error initializing Kuzu database: {e}") + logger.warning("Continuing without database integration") try: await connect_to_jetstream( @@ -735,7 +802,8 @@ async def main(): args.dids_file, args.jetstream_host, thread_depth=args.thread_depth, - comind=comind + comind=comind, + db_manager=db_manager ) except KeyboardInterrupt: logger.info("Shutting down") diff --git a/test_db.py b/test_db.py new file mode 100644 index 0000000..45a4e77 --- /dev/null +++ b/test_db.py @@ -0,0 +1,207 @@ +from src.db_manager import DBManager +import logging +from datetime import datetime +import sys +import json + +# Configure logging to only show errors +logging.basicConfig(level=logging.ERROR) + +def print_header(text): + """Print a header with decoration""" + print(f"\n{'=' * 40}") + print(f" {text}") + print(f"{'=' * 40}\n") + +def print_section(text): + """Print a section header""" + print(f"\n--- {text} ---") + +def check_value(name, value, expected): + """Check if a value matches the expected value and print result""" + if value == expected: + print(f"✅ {name}: PASS") + else: + print(f"❌ {name}: FAIL") + print(f" Expected: {expected}") + print(f" Got: {value}") + +def main(): + print_header("Testing Kuzu Database Integration") + + # Initialize the database manager + print("Initializing database...") + db_manager = DBManager('./demo_db') + print("Database manager initialized.") + + # Test schema creation + print_section("Schema Setup") + try: + db_manager.setup_schema() + print("Schema setup complete.") + except Exception as e: + if "already exists" in str(e): + print("Schema already exists, continuing with tests.") + else: + print(f"Warning: Schema setup encountered an issue: {e}") + print("Continuing with tests anyway...") + + # Test record creation + print_section("Creating Test Records") + test_user_did = "did:test:user123" + current_time = datetime.now().isoformat() + + # Test storing a user + try: + db_manager.store_user( + did=test_user_did, + handle="testuser", + display_name="Test User", + description="This is a test user" + ) + print("✓ Created test user") + except Exception as e: + print(f"✗ Failed to create user: {e}") + + # Test storing a sphere + sphere_uri = f"at://{test_user_did}/me.comind.sphere.core/test-sphere" + try: + db_manager.store_record( + collection="me.comind.sphere.core", + record={ + "title": "Test Sphere", + "text": "This is a test sphere.", + "description": "A sphere for testing", + "createdAt": current_time + }, + uri=sphere_uri, + cid="testcid123", + author_did=test_user_did, + rkey="test-sphere" + ) + print("✓ Created test sphere") + except Exception as e: + print(f"✗ Failed to create sphere: {e}") + + # Test storing a concept + concept_uri = f"at://{test_user_did}/me.comind.blip.concept/test-concept" + try: + db_manager.store_record( + collection="me.comind.blip.concept", + record={ + "text": "test concept", + "createdAt": current_time + }, + uri=concept_uri, + cid="testcid456", + author_did=test_user_did, + rkey="test-concept", + sphere_uri=sphere_uri + ) + print("✓ Created test concept") + except Exception as e: + print(f"✗ Failed to create concept: {e}") + + # Test storing a relationship + link_uri = f"at://{test_user_did}/me.comind.relationship.link/test-link" + try: + db_manager.store_record( + collection="me.comind.relationship.link", + record={ + "relationship": "associated", + "target": concept_uri, + "strength": 0.8, + "note": "Test relation", + "createdAt": current_time + }, + uri=link_uri, + cid="testcid789", + author_did=test_user_did, + rkey="test-link" + ) + print("✓ Created test relationship") + except Exception as e: + print(f"✗ Failed to create relationship: {e}") + + # Test record retrieval + print_section("Record Retrieval Tests") + + # Get sphere record by rkey + try: + sphere = db_manager.get_record("me.comind.sphere.core", "test-sphere") + sphere_found = sphere is not None + check_value("Get sphere by rkey", sphere_found, True) + if sphere: + check_value("Sphere title", sphere.get("title"), "Test Sphere") + except Exception as e: + print(f"✗ Error retrieving sphere by rkey: {e}") + + # Get sphere record by URI + try: + sphere_by_uri = db_manager.get_record_by_uri(sphere_uri) + sphere_by_uri_found = sphere_by_uri is not None + check_value("Get sphere by URI", sphere_by_uri_found, True) + except Exception as e: + print(f"✗ Error retrieving sphere by URI: {e}") + + # List records from collection + try: + concepts = db_manager.list_records("me.comind.blip.concept") + has_concepts = len(concepts) >= 1 + check_value("List concepts collection", has_concepts, True) + if has_concepts: + print(f" Found {len(concepts)} concepts") + except Exception as e: + print(f"✗ Error listing concepts: {e}") + + # Test text search using fallback + print_section("Text Search Tests") + try: + search_results = db_manager._find_records_basic("test", None, 10) + basic_search_found = len(search_results) >= 1 + check_value("Basic search found results", basic_search_found, True) + if basic_search_found: + print(f" Found {len(search_results)} records with basic search") + except Exception as e: + print(f"✗ Error in basic search: {e}") + + # Test full-text search (may fail if FTS extension is not available) + # try: + # fts_results = db_manager.find_similar_records("test") + # fts_found = len(fts_results) >= 1 + # check_value("Full-text search found results", fts_found, True) + # if fts_found: + # print(f" Found {len(fts_results)} records with full-text search") + # except Exception as e: + # print(f"ℹ️ Full-text search test: SKIPPED - {str(e)}") + + # Test relationship queries + print_section("Relationship Tests") + try: + relationships = db_manager.query_relationships(sphere_uri) + has_relationships = len(relationships) >= 1 + check_value("Relationship query returned results", has_relationships, True) + + if has_relationships: + print(f" Found {len(relationships)} relationships from sphere") + rel = relationships[0] + check_value("Relationship source correct", rel.get('source_uri'), sphere_uri) + except Exception as e: + print(f"✗ Error in relationship query: {e}") + + # Test relationship query with specific type + try: + typed_relationships = db_manager.query_relationships(sphere_uri, rel_type="associated") + print(f" Found {len(typed_relationships)} 'associated' relationships") + except Exception as e: + print(f"✗ Error in typed relationship query: {e}") + + print_header("Tests Completed") + +if __name__ == "__main__": + # Capture stdout to see if it's being redirected + try: + main() + except Exception as e: + print(f"ERROR: Test failed with exception: {e}") + sys.exit(1) \ No newline at end of file diff --git a/test_kuzu.py b/test_kuzu.py new file mode 100644 index 0000000..52dbca7 --- /dev/null +++ b/test_kuzu.py @@ -0,0 +1,165 @@ +from src.db_manager import DBManager +import logging +import json +import sys + +# Configure logging to show debug messages +logging.basicConfig(level=logging.INFO) + +def print_header(text): + """Print a header with decoration""" + print(f"\n{'=' * 40}") + print(f" {text}") + print(f"{'=' * 40}\n") + +def print_section(text): + """Print a section header""" + print(f"\n--- {text} ---") + +def main(): + print_header("Testing Kuzu Database ATProto Integration") + + # Initialize the database manager + print("Initializing database...") + db_path = "./demo_db" + db_manager = DBManager(db_path) + print("Database manager initialized.") + + # Set up the schema explicitly + print("Setting up database schema...") + try: + db_manager.setup_schema() + print("Schema setup complete.") + except Exception as e: + if "already exists" in str(e): + print("Schema already exists, continuing with tests.") + else: + print(f"Error setting up schema: {e}") + print("Continuing anyway to see what works...") + + # Create a test record + print_section("Creating Test ATProto Record") + try: + test_did = "did:test:user123" + test_rkey = "testatrecord123" + test_uri = f"at://{test_did}/app.bsky.feed.post/{test_rkey}" + + # Store a test repo + db_manager.conn.execute(""" + MERGE (r:Repo {did: $did}) + SET r.receivedAt = TIMESTAMP('2025-05-14T22:00:00Z') + """, {'did': test_did}) + print("✓ Created test repo") + + # Store a test ATRecord + test_record = { + "text": "This is a test post about AI and KuzuDB integration", + "createdAt": "2025-05-14T22:00:00Z" + } + + result = db_manager.store_atproto_record( + uri=test_uri, + cid="testcid123", + nsid="app.bsky.feed.post", + record=test_record, + author_did=test_did, + rkey=test_rkey, + labels=["ATRecord", "Post", "Test"] + ) + + print(f"✓ Created test ATProto record: {result}") + except Exception as e: + print(f"Error creating test record: {e}") + for i, record in enumerate(records): + uri = record.get('uri', 'Unknown URI') + nsid = record.get('nsid', 'Unknown NSID') + labels = record.get('labels', []) + + print(f"{i+1}. {nsid} - {uri}") + print(f" Labels: {', '.join(labels)}") + print() + except Exception as e: + print(f"Error listing ATProto records: {e}") + + # Count records by NSID + print_section("Record Counts by NSID") + try: + nsids = [ + "app.bsky.feed.post", + "me.comind.blip.concept", + "me.comind.blip.thought", + "me.comind.blip.emotion", + "me.comind.sphere.core" + ] + + for nsid in nsids: + records = db_manager.list_atproto_records(nsid=nsid, limit=1000) + print(f"{nsid}: {len(records)} records") + except Exception as e: + print(f"Error counting records: {e}") + + # Query relationships + print_section("Relationships") + try: + # Get the first post record to use as a source + post_records = db_manager.list_atproto_records(nsid="app.bsky.feed.post", limit=1) + + if post_records: + source_uri = post_records[0]['uri'] + print(f"Querying relationships for {source_uri}") + + relationships = db_manager.query_atproto_relationships(source_uri, limit=10) + print(f"Found {len(relationships)} relationships") + + for i, rel in enumerate(relationships): + rel_type = rel.get('type', 'Unknown') + target_uri = rel.get('target_uri', 'Unknown target') + + print(f"{i+1}. {rel_type} -> {target_uri}") + else: + print("No post records found to query relationships") + except Exception as e: + print(f"Error querying relationships: {e}") + + # Test searching + print_section("Text Search") + try: + search_text = "AI" + print(f"Searching for '{search_text}'...") + + # Check if db_manager has the find_similar_records method + if not hasattr(db_manager, 'find_similar_records'): + print("find_similar_records method not available, using _find_records_basic") + results = db_manager._find_records_basic(search_text, limit=5) + else: + try: + # Try full-text search first + results = db_manager.find_similar_records(search_text, limit=5) + except Exception as search_error: + print(f"Full-text search failed: {search_error}") + print("Falling back to basic search...") + results = db_manager._find_records_basic(search_text, limit=5) + + print(f"Search found {len(results)} results") + + for i, result in enumerate(results): + uri = result.get('uri', 'Unknown URI') + collection = result.get('collection', 'Unknown collection') + value = result.get('value', {}) + + content_preview = str(value)[:100] + "..." if len(str(value)) > 100 else str(value) + print(f"{i+1}. {collection} - {uri}") + print(f" Content: {content_preview}") + print() + except Exception as e: + print(f"Error during search: {e}") + print("Skipping search test due to error") + + print_header("Tests Completed") + +if __name__ == "__main__": + try: + main() + except Exception as e: + print(f"ERROR: Test failed with exception: {e}") + sys.exit(1) \ No newline at end of file -- 2.51.2