docuquery 0.1.0__py3-none-any.whl

docuquery/main.py

@@ -0,0 +1,157 @@
+import logging
+import os
+import uuid
+
+import click
+
+
+@click.group()
+@click.version_option(package_name="docuquery")
+def cli() -> None:
+    from docuquery.utils import ensure_env_variables
+
+    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+    os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+    # PyTorch warns on macOS that a distributed-training feature is missing.
+    # DocuQuery does not use it, so hide the warning.
+    logging.getLogger("torch.distributed.elastic.multiprocessing.redirects").setLevel(
+        logging.ERROR
+    )
+    ensure_env_variables()
+    pass
+
+
+@click.command(name="init")
+def init() -> None:
+    """Initializes the local environment: creates tables in the Postgres database."""
+    # Imported here: constants.py reads the environment on import, so it must
+    # run after ensure_env_variables() has loaded the config file.
+    from docuquery.database import Database
+
+    try:
+        click.echo("Initializing database...", color=True)
+        db = Database()
+        db.initialize_database()
+        click.secho("✔ Database initialized successfully.", fg="green")
+    except click.ClickException:
+        raise
+    except Exception as e:
+        raise click.ClickException(f"Could not initialize database: {e}") from e
+
+
+@click.command(name="ingest")
+@click.argument(
+    "file_path",
+    type=click.Path(exists=True),
+    required=True,
+    help="The path to the document file to be ingested. You can drag and drop the file here.",
+)
+def ingest(file_path: str) -> None:
+    """Ingest a document from the specified file path."""
+    from docuquery.ingest import DocumentIngestor
+
+    try:
+        ingestor = DocumentIngestor()
+        ingestor.process_and_store_document(file_path)
+    except click.ClickException:
+        raise
+    except Exception as e:
+        raise click.ClickException(f"Could not ingest the document: {e}") from e
+
+
+@click.command(name="ask")
+@click.argument(
+    "question", required=False, help="Start the conversation with this question."
+)
+@click.option(
+    "--session-id", required=False, help="Resume a saved conversation by session ID."
+)
+def ask(question: str, session_id: str) -> None:
+    """Chat with your documents. Type 'exit' or press Ctrl-D to quit."""
+    from docuquery.rag import RAGBase
+
+    rag = RAGBase()
+
+    if session_id:
+        try:
+            uuid.UUID(session_id)
+        except ValueError:
+            raise click.ClickException(
+                "The session could not be loaded. Invalid session ID."
+            ) from None
+        try:
+            session = rag.load_conversation(session_id)
+            click.echo(f'Resumed session: {session_id} - Title: "{session["title"]}"')
+        except click.ClickException:
+            raise
+        except ValueError:
+            raise click.ClickException(
+                "The session could not be loaded. Session does not exist."
+            ) from None
+
+    title = question
+    try:
+        while True:
+            if not question:
+                question = click.prompt(
+                    click.style("You", fg="cyan"), prompt_suffix=": "
+                ).strip()
+            if question.lower() in ("exit", "quit"):
+                break
+            if question:
+                title = title or question
+                click.secho("Answer: ", nl=False, fg="green")
+                try:
+                    for token in rag.query_rag_stream(question):
+                        click.secho(token, nl=False)
+                    # Save conversation
+                    session_id = rag.save_conversation(title=title[:50])
+                except Exception as e:
+                    click.secho(
+                        f"❌ Error: Could not generate an answer: {e}",
+                        err=True,
+                        fg="red",
+                    )
+                click.echo("\n")
+            question = None
+    except (KeyboardInterrupt, EOFError, click.Abort):
+        pass
+    click.echo("\nGoodbye!")
+
+
+@click.command(name="sessions")
+def list_sessions() -> None:
+    """List all saved conversation sessions."""
+    from docuquery.rag import RAGBase
+
+    rag = RAGBase()
+    try:
+        sessions = rag.list_sessions()
+        if not sessions:
+            click.echo("No sessions found.")
+        for session in sessions:
+            click.echo(
+                f'Session ID: {session["id"]} - Title: "{session.get("title", "Untitled")}"'
+            )
+    except click.ClickException:
+        raise
+    except Exception as e:
+        raise click.ClickException(f"Could not list sessions: {e}") from e
+
+
+@click.command(name="resume")
+@click.argument("session_id", required=True, help="The session ID to resume.")
+@click.pass_context
+def resume_session(ctx: click.Context, session_id: str) -> None:
+    """Resume a saved conversation by session ID."""
+    ctx.invoke(ask, question=None, session_id=session_id)
+
+
+cli.add_command(init)
+cli.add_command(ingest)
+cli.add_command(ask)
+cli.add_command(list_sessions)
+cli.add_command(resume_session)
+
+if __name__ == "__main__":
+    cli()

docuquery/rag.py

@@ -0,0 +1,279 @@
+import json
+import typing
+from functools import cached_property
+
+from openai import OpenAI
+from sentence_transformers import SentenceTransformer
+
+from docuquery.constants import (
+    DOCUQUERY_OPENAI_MODEL,
+    OPENAI_API_KEY,
+    SENTENCE_TRANSFORMER_MODEL,
+    SYSTEM_INSTRUCTIONS,
+    USER_PROMPT_TEMPLATE,
+)
+from docuquery.database import Database
+from docuquery.tools import TOOLS
+
+
+class RAGBase:
+    MAX_SEARCH_RESULTS = 20
+
+    def __init__(
+        self,
+        instructions: str = SYSTEM_INSTRUCTIONS,
+        prompt_template: str = USER_PROMPT_TEMPLATE,
+        model: str = DOCUQUERY_OPENAI_MODEL,
+    ) -> None:
+        self.db = Database()
+        self.openai_client = OpenAI(api_key=OPENAI_API_KEY)
+
+        self.instructions = instructions
+        self.prompt_template = prompt_template
+        self.model = model
+
+        self.messages = [{"role": "system", "content": self.instructions}]
+        self.saved_count = 1
+        self.session_id = None
+
+    @cached_property
+    def embedder(self) -> SentenceTransformer:
+        return SentenceTransformer(SENTENCE_TRANSFORMER_MODEL)
+
+    def save_conversation(self, title: str) -> str:
+        self.session_id = self.__start_new_session(title)
+        self.db.add_messages(
+            conversation_id=self.session_id, messages=self.messages[self.saved_count :]
+        )
+        self.saved_count = len(self.messages)
+        return self.session_id
+
+    def load_conversation(self, session_id: str) -> dict:
+        """Load a conversation session by its session ID.
+
+        Args:
+            session_id (str): The ID of the conversation session to load.
+        Raises:
+            ValueError: If no conversation is found for the given session ID.
+        Returns:
+            dict: The conversation session data, including its messages.
+        """
+        conversation = self.db.get_conversation_by_id(conversation_id=session_id)
+        if conversation is None:
+            raise ValueError(f"No conversation found for session_id: {session_id}")
+        message_content = self.instructions + "\n\n Session ID: " + session_id
+        self.messages = [{"role": "system", "content": message_content}]
+        self.messages.extend(conversation["messages"])
+        self.session_id = session_id
+        self.saved_count = len(self.messages)
+        return conversation
+
+    def list_sessions(self, page: int = 1, limit: int = 20) -> list[dict]:
+        """Get the list of conversation sessions
+
+        Args:
+            page (int, optional): The page number to retrieve. Defaults to 1.
+            limit (int, optional): The number of sessions per page. Defaults to 20.
+
+        Returns:
+            list[dict]: Returns a list of conversation sessions, each represented as a dictionary with 'id' and 'title' keys.
+        """
+        return self.db.list_conversations(page=page, page_size=limit)
+
+    def search(self, query: str, limit: int = 3) -> list[dict]:
+        limit = (
+            1
+            if limit < 1
+            else self.MAX_SEARCH_RESULTS
+            if limit > self.MAX_SEARCH_RESULTS
+            else limit
+        )
+        query_embedding = self.embedder.encode(query)
+        return self.db.search_similar_chunks(
+            query_embedding=query_embedding, limit=limit
+        )
+
+    def get_session_messages(self, session_id: str) -> str:
+        """Retrieve the formatted conversation messages for a given session ID.
+
+        Args:
+            session_id (str): The ID of the conversation session to load.
+
+        Returns:
+            str: The formatted conversation messages for the session, or an error message if the session does not exist.
+        """
+        conversation = self.db.get_conversation_by_id(conversation_id=session_id)
+        if conversation is None:
+            return "Session does not exist."
+        return "\n".join(
+            f"{m['role']}: {m.get('content') or '[tool call]'}"
+            for m in conversation["messages"]
+            if m["role"] in ("user", "assistant")
+        )
+
+    def build_context(self, matches: list[dict]) -> str:
+
+        # Loop through the matches and convert the dictionary into formated strings assuming the keys are unknown.
+        context_str = ""
+
+        if not matches:
+            return "No relevant context found."
+
+        for match in matches:
+            for key in match:
+                if match[key] is None:
+                    continue
+                field = key.split("_")
+                field = " ".join([word.capitalize() for word in field])
+                match[key] = f"{field}: {match[key]}"
+                context_str += f"{match[key]}\n"
+
+            context_str += "\n\n---\n\n"
+
+        return context_str
+
+    def build_prompt(self, matches: list[dict]) -> str:
+        context = self.build_context(matches)
+        return self.prompt_template.format(context=context)
+
+    def __run_tool(self, name: str, args: dict) -> str:
+        if name == "search":
+            matches = self.search(args["query"], limit=args.get("limit", 3))
+            return self.build_context(matches)
+        if name == "__count_documents":
+            return str(self.__count_documents(**args))
+        if name == "__list_documents":
+            matches = self.__list_documents(**args)
+            return self.build_context(matches)
+        if name == "list_sessions":
+            sessions = self.list_sessions(**args)
+            return self.build_context(sessions)
+        if name == "get_session_messages":
+            return self.get_session_messages(**args)
+        return f"Unknown tool: {name}"
+
+    def query_rag_stream(
+        self, prompt: str, max_turns: int = 3
+    ) -> typing.Generator[str, None, None]:
+        self.messages.append({"role": "user", "content": prompt})
+
+        for _ in range(max_turns):
+            # Stream response tokens from OpenAI
+            stream = self.openai_client.chat.completions.create(
+                model=self.model,
+                messages=self.messages,
+                stream=True,
+                temperature=0.0,
+                tools=TOOLS,
+            )
+
+            text_parts = []
+            tool_calls = {}
+            finish_reason = None
+
+            for chunk in stream:
+                if not chunk.choices:
+                    continue
+
+                choice = chunk.choices[0]
+                delta = choice.delta
+                if delta.content:
+                    text_parts.append(delta.content)
+                    yield delta.content
+
+                for tc in delta.tool_calls or []:
+                    call = tool_calls.setdefault(
+                        tc.index, {"id": None, "name": "", "arguments": ""}
+                    )
+                    if tc.id:
+                        call["id"] = tc.id
+                    if tc.function and tc.function.name:
+                        call["name"] += tc.function.name
+                    if tc.function and tc.function.arguments:
+                        call["arguments"] += tc.function.arguments
+
+                if choice.finish_reason:
+                    finish_reason = choice.finish_reason
+
+            if finish_reason != "tool_calls":
+                self.messages.append(
+                    {"role": "assistant", "content": "".join(text_parts) or None}
+                )
+                return
+
+            calls = [tool_calls[i] for i in sorted(tool_calls)]
+            self.messages.append(
+                {
+                    "role": "assistant",
+                    "content": "".join(text_parts) or None,
+                    "tool_calls": [
+                        {
+                            "id": c["id"],
+                            "type": "function",
+                            "function": {
+                                "name": c["name"],
+                                "arguments": c["arguments"],
+                            },
+                        }
+                        for c in calls
+                    ],
+                }
+            )
+
+            for call in calls:
+                function_name = call["name"]
+                try:
+                    args = json.loads(call["arguments"] or "{}")
+                    result = self.__run_tool(function_name, args)
+                except Exception as e:
+                    result = (
+                        f"Tool error: {e}"  # let the model see the failure and recover
+                    )
+                self.messages.append(
+                    {"role": "tool", "tool_call_id": call["id"], "content": result}
+                )
+
+        yield "\n[Stopped: reached max tool-call turns]"
+
+    def __count_documents(self, query: str) -> int:
+        """Searches for documents, and counts the total number of documents that match the query.
+
+        Args:
+            query (str): The search query to match against document names, content, and embeddings.
+
+        Returns:
+            int: The total number of documents that match the query.
+        """
+        filters = {}
+        if query.strip():
+            filters["query_embedding"] = self.embedder.encode(query)
+        return self.db.count_documents(**filters)
+
+    def __list_documents(
+        self, query: str = "", page: int = 1, limit: int = 5
+    ) -> list[dict]:
+        """Lists documents, optionally filtering by a query string and supporting pagination.
+
+        Args:
+            query (str): Optional search query to filter document names and content.
+            limit (int): Maximum number of documents to return per page.
+            page (int): The page number for pagination.
+
+        Returns:
+            list[dict]: A list of dictionaries containing document information.
+        """
+        filters = {}
+        if query.strip():
+            filters["query_embedding"] = self.embedder.encode(query)
+        return self.db.list_documents(
+            page=page,
+            page_size=limit,
+            **filters,
+        )
+
+    def __start_new_session(self, title: str) -> str:
+        if self.session_id:
+            return self.session_id
+
+        self.session_id = self.db.create_conversation(title=title)
+        return self.session_id

docuquery/tools.py

@@ -0,0 +1,118 @@
+search_tool = {
+    "type": "function",
+    "function": {
+        "name": "search",
+        "description": "Search the user's ingested documents for passages relevant to a query.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The search query to match against document content.",
+                },
+                "limit": {
+                    "type": "number",
+                    "description": "Maximum number of search results to return.",
+                },
+            },
+            "required": ["query"],
+            "additionalProperties": False,
+        },
+    },
+}
+
+count_documents_tool = {
+    "type": "function",
+    "function": {
+        "name": "__count_documents",
+        "description": "Searches for documents, and counts the total number of documents that match the query.\n\nReturns:\n    int: The total number of documents that match thequery.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The search query to match against document names, content, and embeddings.",
+                }
+            },
+            "required": ["query"],
+            "additionalProperties": False,
+        },
+    },
+}
+
+list_documents_tool = {
+    "type": "function",
+    "function": {
+        "name": "__list_documents",
+        "description": "Lists documents, optionally filtering by a query string and supporting pagination.\n\nReturns:\n    list[dict]: A list of dictionaries containing document information.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Optional search query to filter document names and content.",
+                },
+                "page": {
+                    "type": "number",
+                    "description": "The page number for pagination.",
+                },
+                "limit": {
+                    "type": "number",
+                    "description": "Maximum number of documents to return per page.",
+                },
+            },
+            "required": [],
+            "additionalProperties": False,
+        },
+    },
+}
+
+list_sessions_tool = {
+    "type": "function",
+    "function": {
+        "name": "list_sessions",
+        "description": "Get the list of conversation sessions.\n\nReturns:\n    list[dict]: Returns a list of conversation sessions, each represented as a dictionary with 'id' and 'title' keys.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "page": {
+                    "type": "number",
+                    "description": "The page number to retrieve. Defaults to 1.",
+                },
+                "limit": {
+                    "type": "number",
+                    "description": "Maximum number of sessions to return per page. Defaults to 20.",
+                },
+            },
+            "required": [],
+            "additionalProperties": False,
+        },
+    },
+}
+
+get_session_messages_tool = {
+    "type": "function",
+    "function": {
+        "name": "get_session_messages",
+        "description": "Retrieve the formatted conversation messages for a given session ID.\n\nReturns:\n    str: The formatted conversation messages for the session, or an error message if the session does not exist.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "session_id": {
+                    "type": "string",
+                    "description": "The ID of the conversation session to load.",
+                }
+            },
+            "required": ["session_id"],
+            "additionalProperties": False,
+        },
+    },
+}
+
+TOOLS = [
+    search_tool,
+    count_documents_tool,
+    list_documents_tool,
+    list_sessions_tool,
+    get_session_messages_tool,
+]

docuquery/utils.py

@@ -0,0 +1,121 @@
+import os
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import click
+
+
+@dataclass
+class ConfigVariable:
+    key: str
+    name: str
+    desc: str
+    default: str = ""
+    secret: bool = False
+    required: bool = False
+
+
+CONFIG_FILE = Path(click.get_app_dir("docuquery")) / "config.env"
+
+CONFIG_VARIABLES: list[ConfigVariable] = [
+    ConfigVariable(
+        key="DOCUQUERY_POSTGRES_URL",
+        name="Postgres URL",
+        desc="Set the Postgres URL for the Docuquery database",
+        default="postgresql://postgres:postgres@localhost:5432/docuquery_db",
+        required=True,
+    ),
+    ConfigVariable(
+        key="OPENAI_API_KEY",
+        name="OpenAI API Key",
+        desc="Get your OpenAI API Key from https://platform.openai.com/account/api-keys",
+        default="",
+        secret=True,
+        required=True,
+    ),
+    ConfigVariable(
+        key="DOCUQUERY_OPENAI_MODEL",
+        name="OpenAI Model",
+        desc="See list of models you can use: https://platform.openai.com/docs/models",
+        default="gpt-5.4-mini",
+        required=False,
+    ),
+    ConfigVariable(
+        key="SENTENCE_TRANSFORMER_MODEL",
+        name="Sentence Transformer Model",
+        desc="See list of models you can use: https://sbert.net/docs/sentence_transformer/pretrained_models.html#original-models",
+        default="sentence-transformers/all-MiniLM-L6-v2",
+        required=False,
+    ),
+    ConfigVariable(
+        key="DOCUQUERY_CHUNK_SIZE",
+        name="Chunk Size",
+        desc="Set the chunk size for document processing",
+        default="512",
+        required=False,
+    ),
+    ConfigVariable(
+        key="DOCUQUERY_OVERLAP",
+        name="Overlap",
+        desc="Set the overlap size for document processing",
+        default="50",
+        required=False,
+    ),
+]
+
+
+def check_env_variable(name: str) -> bool:
+    """Checks if an environment variable is set."""
+    return bool(os.getenv(name, ""))
+
+
+def set_env_variable(name: str, value: str) -> None:
+    """Sets an environment variable."""
+    os.environ[name] = value
+    CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
+
+    with CONFIG_FILE.open("a") as f:
+        f.write(f"{name}={value}\n")
+    CONFIG_FILE.chmod(0o600)
+
+
+def check_config_file() -> bool:
+    """Checks if the config file exists and loads environment variables from it."""
+    if CONFIG_FILE.exists():
+        for line in CONFIG_FILE.read_text().splitlines():
+            key, _, value = line.partition("=")
+            os.environ.setdefault(key.strip(), value.strip())
+        return True
+    return False
+
+
+def ensure_env_variables() -> None:
+    if "--help" in sys.argv:
+        return
+
+    check_config_file()
+
+    for config in CONFIG_VARIABLES:
+        if check_env_variable(config.key):
+            continue
+
+        click.secho(
+            f"Missing{' required' if config.required else ''} config: {config.name}",
+            fg=f"{'red' if config.required else 'yellow'}",
+        )
+        if config.desc:
+            click.secho(config.desc, err=True, fg="blue")
+
+        while True:
+            value = click.prompt(
+                click.style(f"◉ {config.name}", fg="cyan"),
+                hide_input=config.secret or False,
+                prompt_suffix=" > ",
+                type=str,
+                default=config.default if not config.secret else None,
+            )
+
+            if value:
+                set_env_variable(config.key, value)
+                break