mirror of
https://github.com/james-m-jordan/morphik-core.git
synced 2025-05-09 19:32:38 +00:00
129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
DataBridge interactive CLI.
|
|
Assumes a DataBridge server is running.
|
|
|
|
Usage:
|
|
python shell.py <uri>
|
|
Example: python shell.py "http://test_user:token@localhost:8000"
|
|
|
|
This provides the exact same interface as the Python SDK:
|
|
db.ingest_text("content", metadata={...})
|
|
db.ingest_file("path/to/file")
|
|
db.query("what are the key findings?")
|
|
etc...
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add local SDK to path before other imports
|
|
_SDK_PATH = str(Path(__file__).parent / "sdks" / "python")
|
|
if _SDK_PATH not in sys.path:
|
|
sys.path.insert(0, _SDK_PATH)
|
|
|
|
from databridge import DataBridge # noqa: E402
|
|
|
|
|
|
class DB:
|
|
def __init__(self, uri: str):
|
|
"""Initialize DataBridge with URI"""
|
|
# Convert databridge:// to http:// for localhost
|
|
if "localhost" in uri or "127.0.0.1" in uri:
|
|
uri = uri.replace("databridge://", "http://")
|
|
self.uri = uri
|
|
self._client = DataBridge(self.uri, is_local="localhost" in uri or "127.0.0.1" in uri)
|
|
|
|
def ingest_text(self, content: str, metadata: dict = None) -> dict:
|
|
"""Ingest text content into DataBridge"""
|
|
doc = self._client.ingest_text(content, metadata=metadata or {})
|
|
return doc.model_dump()
|
|
|
|
def ingest_file(
|
|
self, file: str, filename: str, metadata: dict = None, content_type: str = None
|
|
) -> dict:
|
|
"""Ingest a file into DataBridge"""
|
|
file_path = Path(file)
|
|
doc = self._client.ingest_file(
|
|
file=file_path, filename=filename, content_type=content_type, metadata=metadata or {}
|
|
)
|
|
return doc.model_dump()
|
|
|
|
def retrieve_chunks(
|
|
self, query: str, filters: dict = None, k: int = 4, min_score: float = 0.0
|
|
) -> list:
|
|
"""Search for relevant chunks"""
|
|
results = self._client.retrieve_chunks(
|
|
query, filters=filters or {}, k=k, min_score=min_score
|
|
)
|
|
return [r.model_dump() for r in results]
|
|
|
|
def retrieve_docs(
|
|
self, query: str, filters: dict = None, k: int = 4, min_score: float = 0.0
|
|
) -> list:
|
|
"""Retrieve relevant documents"""
|
|
results = self._client.retrieve_docs(query, filters=filters or {}, k=k, min_score=min_score)
|
|
return [r.model_dump() for r in results]
|
|
|
|
def query(
|
|
self,
|
|
query: str,
|
|
filters: dict = None,
|
|
k: int = 4,
|
|
min_score: float = 0.0,
|
|
max_tokens: int = None,
|
|
temperature: float = None,
|
|
) -> dict:
|
|
"""Generate completion using relevant chunks as context"""
|
|
response = self._client.query(
|
|
query,
|
|
filters=filters or {},
|
|
k=k,
|
|
min_score=min_score,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
)
|
|
return response.model_dump()
|
|
|
|
def list_documents(self, skip: int = 0, limit: int = 100, filters: dict = None) -> list:
|
|
"""List accessible documents"""
|
|
docs = self._client.list_documents(skip=skip, limit=limit, filters=filters or {})
|
|
return [doc.model_dump() for doc in docs]
|
|
|
|
def get_document(self, document_id: str) -> dict:
|
|
"""Get document metadata by ID"""
|
|
doc = self._client.get_document(document_id)
|
|
return doc.model_dump()
|
|
|
|
def close(self):
|
|
"""Close the client connection"""
|
|
self._client.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 2:
|
|
print("Error: URI argument required")
|
|
print(__doc__)
|
|
sys.exit(1)
|
|
|
|
# Create DB instance with provided URI
|
|
db = DB(sys.argv[1])
|
|
|
|
# Start an interactive Python shell with 'db' already imported
|
|
import code
|
|
import readline # Enable arrow key history
|
|
import rlcompleter # noqa: F401 # Enable tab completion
|
|
|
|
readline.parse_and_bind("tab: complete")
|
|
|
|
# Create the interactive shell
|
|
shell = code.InteractiveConsole(locals())
|
|
|
|
# Print welcome message
|
|
print("\nDataBridge CLI ready to use. The 'db' object is available with all SDK methods.")
|
|
print("Example: db.ingest_text('hello world')")
|
|
print("Type help(db) for documentation.")
|
|
|
|
# Start the shell
|
|
shell.interact(banner="")
|