diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py new file mode 100644 index 0000000..ec4f703 --- /dev/null +++ b/ai-hub/app/db/models.py @@ -0,0 +1,159 @@ +from datetime import datetime +from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON +from sqlalchemy.orm import relationship + +# The declarative_base class is already defined in database.py. +# We will import it from there to ensure all models use the same base. +from .database import Base + +# --- SQLAlchemy Models --- +# These classes define the structure of the database tables and how they relate. + +class Session(Base): + """ + SQLAlchemy model for the 'sessions' table. + + Each session represents a single conversation between a user and the AI. + It links a user to a series of messages. + """ + __tablename__ = 'sessions' + + # Primary key for the session. + id = Column(Integer, primary_key=True, index=True) + # The ID of the user who owns this session. + user_id = Column(String, index=True, nullable=False) + # A title for the conversation, which can be generated by the AI. + title = Column(String, index=True, nullable=True) + # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek"). + model_name = Column(String, nullable=True) + # Timestamp for when the session was created. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # Flag to indicate if the session has been archived or soft-deleted. + is_archived = Column(Boolean, default=False, nullable=False) + + # Defines a one-to-many relationship with the Message table. + # 'back_populates' tells SQLAlchemy that there's a corresponding relationship + # on the other side. 'cascade' ensures that when a session is deleted, + # all its associated messages are also deleted. + messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Message(Base): + """ + SQLAlchemy model for the 'messages' table. + + This table stores the individual chat messages within a session, + including who sent them (user or AI) and the content. + """ + __tablename__ = 'messages' + + # Primary key for the message. + id = Column(Integer, primary_key=True, index=True) + # The foreign key that links this message to its parent session. + # This is a critical link for reconstructing chat history. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) + # Identifies the sender of the message, e.g., 'user' or 'assistant'. + sender = Column(String, nullable=False) + # The actual text content of the message. + content = Column(Text, nullable=False) + # Timestamp for when the message was sent. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # The time taken for the model to generate the response, in seconds. + model_response_time = Column(Integer, nullable=True) + # The number of tokens in the message (both input and output). + token_count = Column(Integer, nullable=True) + # A JSON field to store unstructured metadata about the message, such as tool calls. + # This column has been renamed from 'metadata' to avoid a conflict. + message_metadata = Column(JSON, nullable=True) + + + # Relationship back to the parent Session. + # This allows us to access the parent Session object from a Message object. + session = relationship("Session", back_populates="messages") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Document(Base): + """ + SQLAlchemy model for the 'documents' table. + + This table stores the metadata and original text content of a document. + The content is the data that will be chunked, embedded, and used for RAG. + """ + __tablename__ = 'documents' + + # Primary key for the document, uniquely identifying each entry. + id = Column(Integer, primary_key=True, index=True) + # The title of the document for easy human-readable reference. + title = Column(String, index=True, nullable=False) + # The actual text content of the document. Using Text for potentially long strings. + text = Column(Text, nullable=False) + # The original source URL or path of the document. + source_url = Column(String, nullable=True) + # A string to identify the author of the document. + author = Column(String, nullable=True) + # The current processing status of the document (e.g., 'ready', 'processing', 'failed'). + status = Column(String, default="processing", nullable=False) + # Timestamp for when the document was added to the database. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # A string to identify the user who added the document, useful for multi-user apps. + user_id = Column(String, index=True, nullable=True) + + # Defines a one-to-one relationship with the VectorMetadata table. + vector_metadata = relationship( + "VectorMetadata", + back_populates="document", + cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted. + uselist=False + ) + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class VectorMetadata(Base): + """ + SQLAlchemy model for the 'vector_metadata' table. + + This table links a document to its corresponding vector representation + in the FAISS index. It is critical for syncing data between the + relational database and the vector store. + """ + __tablename__ = 'vector_metadata' + + # Primary key for the metadata entry. + id = Column(Integer, primary_key=True, index=True) + # Foreign key that links this metadata entry back to its Document. + document_id = Column(Integer, ForeignKey('documents.id'), unique=True) + # The index number in the FAISS vector store where the vector for this document is stored. + faiss_index = Column(Integer, nullable=False, index=True) + # Foreign key to link this vector metadata to a specific session. + # This is crucial for retrieving relevant RAG context for a given conversation. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True) + # The name of the embedding model used to create the vector. + embedding_model = Column(String, nullable=False) + + # Defines a many-to-one relationship with the Document table. + document = relationship("Document", back_populates="vector_metadata") + # Defines a many-to-one relationship with the Session table. + session = relationship("Session") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py new file mode 100644 index 0000000..ec4f703 --- /dev/null +++ b/ai-hub/app/db/models.py @@ -0,0 +1,159 @@ +from datetime import datetime +from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON +from sqlalchemy.orm import relationship + +# The declarative_base class is already defined in database.py. +# We will import it from there to ensure all models use the same base. +from .database import Base + +# --- SQLAlchemy Models --- +# These classes define the structure of the database tables and how they relate. + +class Session(Base): + """ + SQLAlchemy model for the 'sessions' table. + + Each session represents a single conversation between a user and the AI. + It links a user to a series of messages. + """ + __tablename__ = 'sessions' + + # Primary key for the session. + id = Column(Integer, primary_key=True, index=True) + # The ID of the user who owns this session. + user_id = Column(String, index=True, nullable=False) + # A title for the conversation, which can be generated by the AI. + title = Column(String, index=True, nullable=True) + # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek"). + model_name = Column(String, nullable=True) + # Timestamp for when the session was created. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # Flag to indicate if the session has been archived or soft-deleted. + is_archived = Column(Boolean, default=False, nullable=False) + + # Defines a one-to-many relationship with the Message table. + # 'back_populates' tells SQLAlchemy that there's a corresponding relationship + # on the other side. 'cascade' ensures that when a session is deleted, + # all its associated messages are also deleted. + messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Message(Base): + """ + SQLAlchemy model for the 'messages' table. + + This table stores the individual chat messages within a session, + including who sent them (user or AI) and the content. + """ + __tablename__ = 'messages' + + # Primary key for the message. + id = Column(Integer, primary_key=True, index=True) + # The foreign key that links this message to its parent session. + # This is a critical link for reconstructing chat history. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) + # Identifies the sender of the message, e.g., 'user' or 'assistant'. + sender = Column(String, nullable=False) + # The actual text content of the message. + content = Column(Text, nullable=False) + # Timestamp for when the message was sent. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # The time taken for the model to generate the response, in seconds. + model_response_time = Column(Integer, nullable=True) + # The number of tokens in the message (both input and output). + token_count = Column(Integer, nullable=True) + # A JSON field to store unstructured metadata about the message, such as tool calls. + # This column has been renamed from 'metadata' to avoid a conflict. + message_metadata = Column(JSON, nullable=True) + + + # Relationship back to the parent Session. + # This allows us to access the parent Session object from a Message object. + session = relationship("Session", back_populates="messages") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Document(Base): + """ + SQLAlchemy model for the 'documents' table. + + This table stores the metadata and original text content of a document. + The content is the data that will be chunked, embedded, and used for RAG. + """ + __tablename__ = 'documents' + + # Primary key for the document, uniquely identifying each entry. + id = Column(Integer, primary_key=True, index=True) + # The title of the document for easy human-readable reference. + title = Column(String, index=True, nullable=False) + # The actual text content of the document. Using Text for potentially long strings. + text = Column(Text, nullable=False) + # The original source URL or path of the document. + source_url = Column(String, nullable=True) + # A string to identify the author of the document. + author = Column(String, nullable=True) + # The current processing status of the document (e.g., 'ready', 'processing', 'failed'). + status = Column(String, default="processing", nullable=False) + # Timestamp for when the document was added to the database. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # A string to identify the user who added the document, useful for multi-user apps. + user_id = Column(String, index=True, nullable=True) + + # Defines a one-to-one relationship with the VectorMetadata table. + vector_metadata = relationship( + "VectorMetadata", + back_populates="document", + cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted. + uselist=False + ) + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class VectorMetadata(Base): + """ + SQLAlchemy model for the 'vector_metadata' table. + + This table links a document to its corresponding vector representation + in the FAISS index. It is critical for syncing data between the + relational database and the vector store. + """ + __tablename__ = 'vector_metadata' + + # Primary key for the metadata entry. + id = Column(Integer, primary_key=True, index=True) + # Foreign key that links this metadata entry back to its Document. + document_id = Column(Integer, ForeignKey('documents.id'), unique=True) + # The index number in the FAISS vector store where the vector for this document is stored. + faiss_index = Column(Integer, nullable=False, index=True) + # Foreign key to link this vector metadata to a specific session. + # This is crucial for retrieving relevant RAG context for a given conversation. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True) + # The name of the embedding model used to create the vector. + embedding_model = Column(String, nullable=False) + + # Defines a many-to-one relationship with the Document table. + document = relationship("Document", back_populates="vector_metadata") + # Defines a many-to-one relationship with the Session table. + session = relationship("Session") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" diff --git a/ai-hub/pyproject.toml b/ai-hub/pyproject.toml new file mode 100644 index 0000000..2071533 --- /dev/null +++ b/ai-hub/pyproject.toml @@ -0,0 +1,5 @@ +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = [ + "tests" +] \ No newline at end of file diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py new file mode 100644 index 0000000..ec4f703 --- /dev/null +++ b/ai-hub/app/db/models.py @@ -0,0 +1,159 @@ +from datetime import datetime +from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON +from sqlalchemy.orm import relationship + +# The declarative_base class is already defined in database.py. +# We will import it from there to ensure all models use the same base. +from .database import Base + +# --- SQLAlchemy Models --- +# These classes define the structure of the database tables and how they relate. + +class Session(Base): + """ + SQLAlchemy model for the 'sessions' table. + + Each session represents a single conversation between a user and the AI. + It links a user to a series of messages. + """ + __tablename__ = 'sessions' + + # Primary key for the session. + id = Column(Integer, primary_key=True, index=True) + # The ID of the user who owns this session. + user_id = Column(String, index=True, nullable=False) + # A title for the conversation, which can be generated by the AI. + title = Column(String, index=True, nullable=True) + # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek"). + model_name = Column(String, nullable=True) + # Timestamp for when the session was created. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # Flag to indicate if the session has been archived or soft-deleted. + is_archived = Column(Boolean, default=False, nullable=False) + + # Defines a one-to-many relationship with the Message table. + # 'back_populates' tells SQLAlchemy that there's a corresponding relationship + # on the other side. 'cascade' ensures that when a session is deleted, + # all its associated messages are also deleted. + messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Message(Base): + """ + SQLAlchemy model for the 'messages' table. + + This table stores the individual chat messages within a session, + including who sent them (user or AI) and the content. + """ + __tablename__ = 'messages' + + # Primary key for the message. + id = Column(Integer, primary_key=True, index=True) + # The foreign key that links this message to its parent session. + # This is a critical link for reconstructing chat history. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) + # Identifies the sender of the message, e.g., 'user' or 'assistant'. + sender = Column(String, nullable=False) + # The actual text content of the message. + content = Column(Text, nullable=False) + # Timestamp for when the message was sent. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # The time taken for the model to generate the response, in seconds. + model_response_time = Column(Integer, nullable=True) + # The number of tokens in the message (both input and output). + token_count = Column(Integer, nullable=True) + # A JSON field to store unstructured metadata about the message, such as tool calls. + # This column has been renamed from 'metadata' to avoid a conflict. + message_metadata = Column(JSON, nullable=True) + + + # Relationship back to the parent Session. + # This allows us to access the parent Session object from a Message object. + session = relationship("Session", back_populates="messages") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Document(Base): + """ + SQLAlchemy model for the 'documents' table. + + This table stores the metadata and original text content of a document. + The content is the data that will be chunked, embedded, and used for RAG. + """ + __tablename__ = 'documents' + + # Primary key for the document, uniquely identifying each entry. + id = Column(Integer, primary_key=True, index=True) + # The title of the document for easy human-readable reference. + title = Column(String, index=True, nullable=False) + # The actual text content of the document. Using Text for potentially long strings. + text = Column(Text, nullable=False) + # The original source URL or path of the document. + source_url = Column(String, nullable=True) + # A string to identify the author of the document. + author = Column(String, nullable=True) + # The current processing status of the document (e.g., 'ready', 'processing', 'failed'). + status = Column(String, default="processing", nullable=False) + # Timestamp for when the document was added to the database. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # A string to identify the user who added the document, useful for multi-user apps. + user_id = Column(String, index=True, nullable=True) + + # Defines a one-to-one relationship with the VectorMetadata table. + vector_metadata = relationship( + "VectorMetadata", + back_populates="document", + cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted. + uselist=False + ) + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class VectorMetadata(Base): + """ + SQLAlchemy model for the 'vector_metadata' table. + + This table links a document to its corresponding vector representation + in the FAISS index. It is critical for syncing data between the + relational database and the vector store. + """ + __tablename__ = 'vector_metadata' + + # Primary key for the metadata entry. + id = Column(Integer, primary_key=True, index=True) + # Foreign key that links this metadata entry back to its Document. + document_id = Column(Integer, ForeignKey('documents.id'), unique=True) + # The index number in the FAISS vector store where the vector for this document is stored. + faiss_index = Column(Integer, nullable=False, index=True) + # Foreign key to link this vector metadata to a specific session. + # This is crucial for retrieving relevant RAG context for a given conversation. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True) + # The name of the embedding model used to create the vector. + embedding_model = Column(String, nullable=False) + + # Defines a many-to-one relationship with the Document table. + document = relationship("Document", back_populates="vector_metadata") + # Defines a many-to-one relationship with the Session table. + session = relationship("Session") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" diff --git a/ai-hub/pyproject.toml b/ai-hub/pyproject.toml new file mode 100644 index 0000000..2071533 --- /dev/null +++ b/ai-hub/pyproject.toml @@ -0,0 +1,5 @@ +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = [ + "tests" +] \ No newline at end of file diff --git a/ai-hub/requirements.txt b/ai-hub/requirements.txt index 4077554..9426b1e 100644 --- a/ai-hub/requirements.txt +++ b/ai-hub/requirements.txt @@ -6,6 +6,8 @@ pytest requests anyio +sqlalchemy +psycopg2 pytest-asyncio pytest-tornasync pytest-trio \ No newline at end of file diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py new file mode 100644 index 0000000..ec4f703 --- /dev/null +++ b/ai-hub/app/db/models.py @@ -0,0 +1,159 @@ +from datetime import datetime +from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON +from sqlalchemy.orm import relationship + +# The declarative_base class is already defined in database.py. +# We will import it from there to ensure all models use the same base. +from .database import Base + +# --- SQLAlchemy Models --- +# These classes define the structure of the database tables and how they relate. + +class Session(Base): + """ + SQLAlchemy model for the 'sessions' table. + + Each session represents a single conversation between a user and the AI. + It links a user to a series of messages. + """ + __tablename__ = 'sessions' + + # Primary key for the session. + id = Column(Integer, primary_key=True, index=True) + # The ID of the user who owns this session. + user_id = Column(String, index=True, nullable=False) + # A title for the conversation, which can be generated by the AI. + title = Column(String, index=True, nullable=True) + # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek"). + model_name = Column(String, nullable=True) + # Timestamp for when the session was created. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # Flag to indicate if the session has been archived or soft-deleted. + is_archived = Column(Boolean, default=False, nullable=False) + + # Defines a one-to-many relationship with the Message table. + # 'back_populates' tells SQLAlchemy that there's a corresponding relationship + # on the other side. 'cascade' ensures that when a session is deleted, + # all its associated messages are also deleted. + messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Message(Base): + """ + SQLAlchemy model for the 'messages' table. + + This table stores the individual chat messages within a session, + including who sent them (user or AI) and the content. + """ + __tablename__ = 'messages' + + # Primary key for the message. + id = Column(Integer, primary_key=True, index=True) + # The foreign key that links this message to its parent session. + # This is a critical link for reconstructing chat history. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) + # Identifies the sender of the message, e.g., 'user' or 'assistant'. + sender = Column(String, nullable=False) + # The actual text content of the message. + content = Column(Text, nullable=False) + # Timestamp for when the message was sent. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # The time taken for the model to generate the response, in seconds. + model_response_time = Column(Integer, nullable=True) + # The number of tokens in the message (both input and output). + token_count = Column(Integer, nullable=True) + # A JSON field to store unstructured metadata about the message, such as tool calls. + # This column has been renamed from 'metadata' to avoid a conflict. + message_metadata = Column(JSON, nullable=True) + + + # Relationship back to the parent Session. + # This allows us to access the parent Session object from a Message object. + session = relationship("Session", back_populates="messages") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Document(Base): + """ + SQLAlchemy model for the 'documents' table. + + This table stores the metadata and original text content of a document. + The content is the data that will be chunked, embedded, and used for RAG. + """ + __tablename__ = 'documents' + + # Primary key for the document, uniquely identifying each entry. + id = Column(Integer, primary_key=True, index=True) + # The title of the document for easy human-readable reference. + title = Column(String, index=True, nullable=False) + # The actual text content of the document. Using Text for potentially long strings. + text = Column(Text, nullable=False) + # The original source URL or path of the document. + source_url = Column(String, nullable=True) + # A string to identify the author of the document. + author = Column(String, nullable=True) + # The current processing status of the document (e.g., 'ready', 'processing', 'failed'). + status = Column(String, default="processing", nullable=False) + # Timestamp for when the document was added to the database. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # A string to identify the user who added the document, useful for multi-user apps. + user_id = Column(String, index=True, nullable=True) + + # Defines a one-to-one relationship with the VectorMetadata table. + vector_metadata = relationship( + "VectorMetadata", + back_populates="document", + cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted. + uselist=False + ) + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class VectorMetadata(Base): + """ + SQLAlchemy model for the 'vector_metadata' table. + + This table links a document to its corresponding vector representation + in the FAISS index. It is critical for syncing data between the + relational database and the vector store. + """ + __tablename__ = 'vector_metadata' + + # Primary key for the metadata entry. + id = Column(Integer, primary_key=True, index=True) + # Foreign key that links this metadata entry back to its Document. + document_id = Column(Integer, ForeignKey('documents.id'), unique=True) + # The index number in the FAISS vector store where the vector for this document is stored. + faiss_index = Column(Integer, nullable=False, index=True) + # Foreign key to link this vector metadata to a specific session. + # This is crucial for retrieving relevant RAG context for a given conversation. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True) + # The name of the embedding model used to create the vector. + embedding_model = Column(String, nullable=False) + + # Defines a many-to-one relationship with the Document table. + document = relationship("Document", back_populates="vector_metadata") + # Defines a many-to-one relationship with the Session table. + session = relationship("Session") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" diff --git a/ai-hub/pyproject.toml b/ai-hub/pyproject.toml new file mode 100644 index 0000000..2071533 --- /dev/null +++ b/ai-hub/pyproject.toml @@ -0,0 +1,5 @@ +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = [ + "tests" +] \ No newline at end of file diff --git a/ai-hub/requirements.txt b/ai-hub/requirements.txt index 4077554..9426b1e 100644 --- a/ai-hub/requirements.txt +++ b/ai-hub/requirements.txt @@ -6,6 +6,8 @@ pytest requests anyio +sqlalchemy +psycopg2 pytest-asyncio pytest-tornasync pytest-trio \ No newline at end of file diff --git a/ai-hub/setup.py b/ai-hub/setup.py new file mode 100644 index 0000000..5154feb --- /dev/null +++ b/ai-hub/setup.py @@ -0,0 +1,28 @@ +from setuptools import setup, find_packages + +# Read the requirements from requirements.txt +with open("requirements.txt") as f: + requirements = f.read().splitlines() + +setup( + name="ai-hub", + version="0.1.0", + description="An AI Model Hub Service with PostgreSQL and FAISS integration.", + author="Jerry Xie", + author_email="axieyangb@google.com", + packages=find_packages(), + install_requires=requirements, + # This is a key part of the setup. It defines the 'entry point' + # for the application, so you can run it with `uvicorn run:app`. + entry_points={ + "console_scripts": [ + "ai-hub-server=app.main:app" + ], + }, + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + ], + python_requires='>=3.11', +) \ No newline at end of file diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py new file mode 100644 index 0000000..ec4f703 --- /dev/null +++ b/ai-hub/app/db/models.py @@ -0,0 +1,159 @@ +from datetime import datetime +from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON +from sqlalchemy.orm import relationship + +# The declarative_base class is already defined in database.py. +# We will import it from there to ensure all models use the same base. +from .database import Base + +# --- SQLAlchemy Models --- +# These classes define the structure of the database tables and how they relate. + +class Session(Base): + """ + SQLAlchemy model for the 'sessions' table. + + Each session represents a single conversation between a user and the AI. + It links a user to a series of messages. + """ + __tablename__ = 'sessions' + + # Primary key for the session. + id = Column(Integer, primary_key=True, index=True) + # The ID of the user who owns this session. + user_id = Column(String, index=True, nullable=False) + # A title for the conversation, which can be generated by the AI. + title = Column(String, index=True, nullable=True) + # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek"). + model_name = Column(String, nullable=True) + # Timestamp for when the session was created. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # Flag to indicate if the session has been archived or soft-deleted. + is_archived = Column(Boolean, default=False, nullable=False) + + # Defines a one-to-many relationship with the Message table. + # 'back_populates' tells SQLAlchemy that there's a corresponding relationship + # on the other side. 'cascade' ensures that when a session is deleted, + # all its associated messages are also deleted. + messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Message(Base): + """ + SQLAlchemy model for the 'messages' table. + + This table stores the individual chat messages within a session, + including who sent them (user or AI) and the content. + """ + __tablename__ = 'messages' + + # Primary key for the message. + id = Column(Integer, primary_key=True, index=True) + # The foreign key that links this message to its parent session. + # This is a critical link for reconstructing chat history. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) + # Identifies the sender of the message, e.g., 'user' or 'assistant'. + sender = Column(String, nullable=False) + # The actual text content of the message. + content = Column(Text, nullable=False) + # Timestamp for when the message was sent. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # The time taken for the model to generate the response, in seconds. + model_response_time = Column(Integer, nullable=True) + # The number of tokens in the message (both input and output). + token_count = Column(Integer, nullable=True) + # A JSON field to store unstructured metadata about the message, such as tool calls. + # This column has been renamed from 'metadata' to avoid a conflict. + message_metadata = Column(JSON, nullable=True) + + + # Relationship back to the parent Session. + # This allows us to access the parent Session object from a Message object. + session = relationship("Session", back_populates="messages") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Document(Base): + """ + SQLAlchemy model for the 'documents' table. + + This table stores the metadata and original text content of a document. + The content is the data that will be chunked, embedded, and used for RAG. + """ + __tablename__ = 'documents' + + # Primary key for the document, uniquely identifying each entry. + id = Column(Integer, primary_key=True, index=True) + # The title of the document for easy human-readable reference. + title = Column(String, index=True, nullable=False) + # The actual text content of the document. Using Text for potentially long strings. + text = Column(Text, nullable=False) + # The original source URL or path of the document. + source_url = Column(String, nullable=True) + # A string to identify the author of the document. + author = Column(String, nullable=True) + # The current processing status of the document (e.g., 'ready', 'processing', 'failed'). + status = Column(String, default="processing", nullable=False) + # Timestamp for when the document was added to the database. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # A string to identify the user who added the document, useful for multi-user apps. + user_id = Column(String, index=True, nullable=True) + + # Defines a one-to-one relationship with the VectorMetadata table. + vector_metadata = relationship( + "VectorMetadata", + back_populates="document", + cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted. + uselist=False + ) + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class VectorMetadata(Base): + """ + SQLAlchemy model for the 'vector_metadata' table. + + This table links a document to its corresponding vector representation + in the FAISS index. It is critical for syncing data between the + relational database and the vector store. + """ + __tablename__ = 'vector_metadata' + + # Primary key for the metadata entry. + id = Column(Integer, primary_key=True, index=True) + # Foreign key that links this metadata entry back to its Document. + document_id = Column(Integer, ForeignKey('documents.id'), unique=True) + # The index number in the FAISS vector store where the vector for this document is stored. + faiss_index = Column(Integer, nullable=False, index=True) + # Foreign key to link this vector metadata to a specific session. + # This is crucial for retrieving relevant RAG context for a given conversation. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True) + # The name of the embedding model used to create the vector. + embedding_model = Column(String, nullable=False) + + # Defines a many-to-one relationship with the Document table. + document = relationship("Document", back_populates="vector_metadata") + # Defines a many-to-one relationship with the Session table. + session = relationship("Session") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" diff --git a/ai-hub/pyproject.toml b/ai-hub/pyproject.toml new file mode 100644 index 0000000..2071533 --- /dev/null +++ b/ai-hub/pyproject.toml @@ -0,0 +1,5 @@ +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = [ + "tests" +] \ No newline at end of file diff --git a/ai-hub/requirements.txt b/ai-hub/requirements.txt index 4077554..9426b1e 100644 --- a/ai-hub/requirements.txt +++ b/ai-hub/requirements.txt @@ -6,6 +6,8 @@ pytest requests anyio +sqlalchemy +psycopg2 pytest-asyncio pytest-tornasync pytest-trio \ No newline at end of file diff --git a/ai-hub/setup.py b/ai-hub/setup.py new file mode 100644 index 0000000..5154feb --- /dev/null +++ b/ai-hub/setup.py @@ -0,0 +1,28 @@ +from setuptools import setup, find_packages + +# Read the requirements from requirements.txt +with open("requirements.txt") as f: + requirements = f.read().splitlines() + +setup( + name="ai-hub", + version="0.1.0", + description="An AI Model Hub Service with PostgreSQL and FAISS integration.", + author="Jerry Xie", + author_email="axieyangb@google.com", + packages=find_packages(), + install_requires=requirements, + # This is a key part of the setup. It defines the 'entry point' + # for the application, so you can run it with `uvicorn run:app`. + entry_points={ + "console_scripts": [ + "ai-hub-server=app.main:app" + ], + }, + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + ], + python_requires='>=3.11', +) \ No newline at end of file diff --git a/ai-hub/tests/db/test_database.py b/ai-hub/tests/db/test_database.py new file mode 100644 index 0000000..4f8921a --- /dev/null +++ b/ai-hub/tests/db/test_database.py @@ -0,0 +1,91 @@ +import os +import pytest +import importlib +from sqlalchemy.orm import Session +from sqlalchemy.exc import ResourceClosedError +from sqlalchemy import text +from unittest.mock import patch + +def test_sqlite_mode_initialization(monkeypatch): + """ + Tests if the database initializes in SQLite mode correctly. + """ + # Arrange: Set environment variable for SQLite mode + monkeypatch.setenv("DB_MODE", "sqlite") + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check if the configuration is correct for SQLite + assert database.DB_MODE == "sqlite" + assert "sqlite:///./ai_hub.db" in database.DATABASE_URL + assert "connect_args" in database.engine_args + assert database.engine_args["connect_args"] == {"check_same_thread": False} + + # Cleanup the created SQLite file after test, if it exists + if os.path.exists("ai_hub.db"): + os.remove("ai_hub.db") + +def test_postgres_mode_initialization(monkeypatch): + """ + Tests if the database initializes in PostgreSQL mode with a custom URL. + """ + # Arrange: Set env vars for PostgreSQL mode and a specific URL + monkeypatch.setenv("DB_MODE", "postgres") + monkeypatch.setenv("DATABASE_URL", "postgresql://test_user:test_password@testhost/test_db") + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check if the configuration is correct for PostgreSQL + assert database.DB_MODE == "postgres" + assert database.DATABASE_URL == "postgresql://test_user:test_password@testhost/test_db" + assert "pool_pre_ping" in database.engine_args + +def test_default_to_postgres_mode(monkeypatch): + """ + Tests if the system defaults to PostgreSQL mode when DB_MODE is not set. + """ + # Arrange: Ensure DB_MODE is not set + monkeypatch.delenv("DB_MODE", raising=False) + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check that it defaulted to postgres + assert database.DB_MODE == "postgres" + assert "postgresql://user:password@localhost/ai_hub_db" in database.DATABASE_URL + +@patch('app.db.database.SessionLocal') +def test_get_db_yields_and_closes_session(mock_session_local, monkeypatch): + """ + Tests if the get_db() dependency function yields a valid, active session + and correctly closes it afterward by mocking the session object. + """ + # Arrange: Get the actual get_db function from the module + from app.db import database + + # Configure the mock session returned by SessionLocal() + mock_session = mock_session_local.return_value + + db_generator = database.get_db() + + # Act + # 1. Get the session object from the generator + db_session_instance = next(db_generator) + + # Assert + # 2. Check that the yielded object is our mock session + assert db_session_instance is mock_session + mock_session.close.assert_not_called() # The session should not be closed yet + + # 3. Exhaust the generator to trigger the 'finally' block + with pytest.raises(StopIteration): + next(db_generator) + + # 4. Assert that the close() method was called exactly once. + mock_session.close.assert_called_once() + diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py new file mode 100644 index 0000000..ec4f703 --- /dev/null +++ b/ai-hub/app/db/models.py @@ -0,0 +1,159 @@ +from datetime import datetime +from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON +from sqlalchemy.orm import relationship + +# The declarative_base class is already defined in database.py. +# We will import it from there to ensure all models use the same base. +from .database import Base + +# --- SQLAlchemy Models --- +# These classes define the structure of the database tables and how they relate. + +class Session(Base): + """ + SQLAlchemy model for the 'sessions' table. + + Each session represents a single conversation between a user and the AI. + It links a user to a series of messages. + """ + __tablename__ = 'sessions' + + # Primary key for the session. + id = Column(Integer, primary_key=True, index=True) + # The ID of the user who owns this session. + user_id = Column(String, index=True, nullable=False) + # A title for the conversation, which can be generated by the AI. + title = Column(String, index=True, nullable=True) + # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek"). + model_name = Column(String, nullable=True) + # Timestamp for when the session was created. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # Flag to indicate if the session has been archived or soft-deleted. + is_archived = Column(Boolean, default=False, nullable=False) + + # Defines a one-to-many relationship with the Message table. + # 'back_populates' tells SQLAlchemy that there's a corresponding relationship + # on the other side. 'cascade' ensures that when a session is deleted, + # all its associated messages are also deleted. + messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Message(Base): + """ + SQLAlchemy model for the 'messages' table. + + This table stores the individual chat messages within a session, + including who sent them (user or AI) and the content. + """ + __tablename__ = 'messages' + + # Primary key for the message. + id = Column(Integer, primary_key=True, index=True) + # The foreign key that links this message to its parent session. + # This is a critical link for reconstructing chat history. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) + # Identifies the sender of the message, e.g., 'user' or 'assistant'. + sender = Column(String, nullable=False) + # The actual text content of the message. + content = Column(Text, nullable=False) + # Timestamp for when the message was sent. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # The time taken for the model to generate the response, in seconds. + model_response_time = Column(Integer, nullable=True) + # The number of tokens in the message (both input and output). + token_count = Column(Integer, nullable=True) + # A JSON field to store unstructured metadata about the message, such as tool calls. + # This column has been renamed from 'metadata' to avoid a conflict. + message_metadata = Column(JSON, nullable=True) + + + # Relationship back to the parent Session. + # This allows us to access the parent Session object from a Message object. + session = relationship("Session", back_populates="messages") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Document(Base): + """ + SQLAlchemy model for the 'documents' table. + + This table stores the metadata and original text content of a document. + The content is the data that will be chunked, embedded, and used for RAG. + """ + __tablename__ = 'documents' + + # Primary key for the document, uniquely identifying each entry. + id = Column(Integer, primary_key=True, index=True) + # The title of the document for easy human-readable reference. + title = Column(String, index=True, nullable=False) + # The actual text content of the document. Using Text for potentially long strings. + text = Column(Text, nullable=False) + # The original source URL or path of the document. + source_url = Column(String, nullable=True) + # A string to identify the author of the document. + author = Column(String, nullable=True) + # The current processing status of the document (e.g., 'ready', 'processing', 'failed'). + status = Column(String, default="processing", nullable=False) + # Timestamp for when the document was added to the database. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # A string to identify the user who added the document, useful for multi-user apps. + user_id = Column(String, index=True, nullable=True) + + # Defines a one-to-one relationship with the VectorMetadata table. + vector_metadata = relationship( + "VectorMetadata", + back_populates="document", + cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted. + uselist=False + ) + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class VectorMetadata(Base): + """ + SQLAlchemy model for the 'vector_metadata' table. + + This table links a document to its corresponding vector representation + in the FAISS index. It is critical for syncing data between the + relational database and the vector store. + """ + __tablename__ = 'vector_metadata' + + # Primary key for the metadata entry. + id = Column(Integer, primary_key=True, index=True) + # Foreign key that links this metadata entry back to its Document. + document_id = Column(Integer, ForeignKey('documents.id'), unique=True) + # The index number in the FAISS vector store where the vector for this document is stored. + faiss_index = Column(Integer, nullable=False, index=True) + # Foreign key to link this vector metadata to a specific session. + # This is crucial for retrieving relevant RAG context for a given conversation. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True) + # The name of the embedding model used to create the vector. + embedding_model = Column(String, nullable=False) + + # Defines a many-to-one relationship with the Document table. + document = relationship("Document", back_populates="vector_metadata") + # Defines a many-to-one relationship with the Session table. + session = relationship("Session") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" diff --git a/ai-hub/pyproject.toml b/ai-hub/pyproject.toml new file mode 100644 index 0000000..2071533 --- /dev/null +++ b/ai-hub/pyproject.toml @@ -0,0 +1,5 @@ +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = [ + "tests" +] \ No newline at end of file diff --git a/ai-hub/requirements.txt b/ai-hub/requirements.txt index 4077554..9426b1e 100644 --- a/ai-hub/requirements.txt +++ b/ai-hub/requirements.txt @@ -6,6 +6,8 @@ pytest requests anyio +sqlalchemy +psycopg2 pytest-asyncio pytest-tornasync pytest-trio \ No newline at end of file diff --git a/ai-hub/setup.py b/ai-hub/setup.py new file mode 100644 index 0000000..5154feb --- /dev/null +++ b/ai-hub/setup.py @@ -0,0 +1,28 @@ +from setuptools import setup, find_packages + +# Read the requirements from requirements.txt +with open("requirements.txt") as f: + requirements = f.read().splitlines() + +setup( + name="ai-hub", + version="0.1.0", + description="An AI Model Hub Service with PostgreSQL and FAISS integration.", + author="Jerry Xie", + author_email="axieyangb@google.com", + packages=find_packages(), + install_requires=requirements, + # This is a key part of the setup. It defines the 'entry point' + # for the application, so you can run it with `uvicorn run:app`. + entry_points={ + "console_scripts": [ + "ai-hub-server=app.main:app" + ], + }, + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + ], + python_requires='>=3.11', +) \ No newline at end of file diff --git a/ai-hub/tests/db/test_database.py b/ai-hub/tests/db/test_database.py new file mode 100644 index 0000000..4f8921a --- /dev/null +++ b/ai-hub/tests/db/test_database.py @@ -0,0 +1,91 @@ +import os +import pytest +import importlib +from sqlalchemy.orm import Session +from sqlalchemy.exc import ResourceClosedError +from sqlalchemy import text +from unittest.mock import patch + +def test_sqlite_mode_initialization(monkeypatch): + """ + Tests if the database initializes in SQLite mode correctly. + """ + # Arrange: Set environment variable for SQLite mode + monkeypatch.setenv("DB_MODE", "sqlite") + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check if the configuration is correct for SQLite + assert database.DB_MODE == "sqlite" + assert "sqlite:///./ai_hub.db" in database.DATABASE_URL + assert "connect_args" in database.engine_args + assert database.engine_args["connect_args"] == {"check_same_thread": False} + + # Cleanup the created SQLite file after test, if it exists + if os.path.exists("ai_hub.db"): + os.remove("ai_hub.db") + +def test_postgres_mode_initialization(monkeypatch): + """ + Tests if the database initializes in PostgreSQL mode with a custom URL. + """ + # Arrange: Set env vars for PostgreSQL mode and a specific URL + monkeypatch.setenv("DB_MODE", "postgres") + monkeypatch.setenv("DATABASE_URL", "postgresql://test_user:test_password@testhost/test_db") + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check if the configuration is correct for PostgreSQL + assert database.DB_MODE == "postgres" + assert database.DATABASE_URL == "postgresql://test_user:test_password@testhost/test_db" + assert "pool_pre_ping" in database.engine_args + +def test_default_to_postgres_mode(monkeypatch): + """ + Tests if the system defaults to PostgreSQL mode when DB_MODE is not set. + """ + # Arrange: Ensure DB_MODE is not set + monkeypatch.delenv("DB_MODE", raising=False) + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check that it defaulted to postgres + assert database.DB_MODE == "postgres" + assert "postgresql://user:password@localhost/ai_hub_db" in database.DATABASE_URL + +@patch('app.db.database.SessionLocal') +def test_get_db_yields_and_closes_session(mock_session_local, monkeypatch): + """ + Tests if the get_db() dependency function yields a valid, active session + and correctly closes it afterward by mocking the session object. + """ + # Arrange: Get the actual get_db function from the module + from app.db import database + + # Configure the mock session returned by SessionLocal() + mock_session = mock_session_local.return_value + + db_generator = database.get_db() + + # Act + # 1. Get the session object from the generator + db_session_instance = next(db_generator) + + # Assert + # 2. Check that the yielded object is our mock session + assert db_session_instance is mock_session + mock_session.close.assert_not_called() # The session should not be closed yet + + # 3. Exhaust the generator to trigger the 'finally' block + with pytest.raises(StopIteration): + next(db_generator) + + # 4. Assert that the close() method was called exactly once. + mock_session.close.assert_called_once() + diff --git a/ai-hub/tests/db/test_models.py b/ai-hub/tests/db/test_models.py new file mode 100644 index 0000000..c2e928b --- /dev/null +++ b/ai-hub/tests/db/test_models.py @@ -0,0 +1,201 @@ +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm.exc import DetachedInstanceError +from app.db.models import Base, Session, Message, Document, VectorMetadata + +# Use an in-memory SQLite database for testing. This is fast and ensures +# each test starts with a clean slate. +SQLALCHEMY_DATABASE_URL = "sqlite:///:memory:" + +# Create a database engine +engine = create_engine( + SQLALCHEMY_DATABASE_URL, + connect_args={"check_same_thread": False} +) + +# Create a configured "Session" class +TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + + +@pytest.fixture(scope="function") +def db_session(): + """ + A pytest fixture that creates a new database session for a test. + It creates all tables at the start of the test and drops them at the end. + This ensures that each test is completely isolated. + """ + # Create all tables in the database + Base.metadata.create_all(bind=engine) + db = TestingSessionLocal() + try: + yield db + finally: + db.close() + # Drop all tables after the test is finished + Base.metadata.drop_all(bind=engine) + + +def test_create_all_tables_with_relationships(db_session): + """ + Tests if all tables are created successfully and can be inspected. + """ + # Verify that all tables from the Base metadata were created + assert "sessions" in Base.metadata.tables + assert "messages" in Base.metadata.tables + assert "documents" in Base.metadata.tables + assert "vector_metadata" in Base.metadata.tables + + # Check for foreign key constraints to ensure relationships are set up + session_table = Base.metadata.tables['sessions'] + message_table = Base.metadata.tables['messages'] + document_table = Base.metadata.tables['documents'] + vector_metadata_table = Base.metadata.tables['vector_metadata'] + + assert len(message_table.foreign_keys) == 1 + assert list(message_table.foreign_keys)[0].column.table.name == 'sessions' + + assert len(vector_metadata_table.foreign_keys) == 2 + fk_columns = [fk.column.table.name for fk in vector_metadata_table.foreign_keys] + assert 'documents' in fk_columns + assert 'sessions' in fk_columns + + +def test_create_and_retrieve_session(db_session): + """ + Tests the creation and retrieval of a Session object. + """ + # Create a new session object + new_session = Session(user_id="test-user-123", title="Test Session", model_name="gemini") + + # Add to session and commit to the database + db_session.add(new_session) + db_session.commit() + db_session.refresh(new_session) + + # Retrieve the session from the database by its ID + retrieved_session = db_session.query(Session).filter(Session.id == new_session.id).first() + + # Assert that the retrieved session matches the original + assert retrieved_session is not None + assert retrieved_session.user_id == "test-user-123" + assert retrieved_session.title == "Test Session" + assert retrieved_session.model_name == "gemini" + + +def test_create_message_with_session_relationship(db_session): + """ + Tests the creation of a Message and its relationship with a Session. + """ + # First, create a session + new_session = Session(user_id="test-user-123") + db_session.add(new_session) + db_session.commit() + db_session.refresh(new_session) + + # Then, create a message linked to that session + new_message = Message( + session_id=new_session.id, + sender="user", + content="Hello, world!", + token_count=3, + model_response_time=0 + ) + db_session.add(new_message) + db_session.commit() + db_session.refresh(new_message) + + # Retrieve the message and verify its relationship + retrieved_message = db_session.query(Message).filter(Message.id == new_message.id).first() + + assert retrieved_message is not None + assert retrieved_message.session_id == new_session.id + assert retrieved_message.content == "Hello, world!" + # Verify the relationship attribute works + assert retrieved_message.session.user_id == "test-user-123" + + +def test_create_document_and_vector_metadata_relationship(db_session): + """ + Tests the creation of a Document and its linked VectorMetadata. + """ + # Create a new document + new_document = Document( + user_id="test-user-123", + title="Sample Doc", + text="This is some sample text.", + status="ready" + ) + db_session.add(new_document) + db_session.commit() + db_session.refresh(new_document) + + # Create vector metadata linked to the document + vector_meta = VectorMetadata( + document_id=new_document.id, + faiss_index=123, + embedding_model="test-model" + ) + db_session.add(vector_meta) + db_session.commit() + db_session.refresh(vector_meta) + + # Retrieve the document and verify the relationship + retrieved_doc = db_session.query(Document).filter(Document.id == new_document.id).first() + assert retrieved_doc is not None + assert retrieved_doc.vector_metadata is not None + assert retrieved_doc.vector_metadata.faiss_index == 123 + + +def test_cascade_delete_session_and_messages(db_session): + """ + Tests that deleting a session automatically deletes its associated messages due to cascading. + """ + # Create a session and some messages + new_session = Session(user_id="cascade-test") + db_session.add(new_session) + db_session.commit() + db_session.refresh(new_session) + + message1 = Message(session_id=new_session.id, sender="user", content="Msg 1") + message2 = Message(session_id=new_session.id, sender="user", content="Msg 2") + db_session.add_all([message1, message2]) + db_session.commit() + + # Check that messages exist before deletion + assert db_session.query(Message).filter(Message.session_id == new_session.id).count() == 2 + + # Delete the session + db_session.delete(new_session) + db_session.commit() + + # Check that the session is gone and the messages have been cascaded + assert db_session.query(Session).filter(Session.id == new_session.id).count() == 0 + assert db_session.query(Message).filter(Message.session_id == new_session.id).count() == 0 + + +def test_cascade_delete_document_and_vector_metadata(db_session): + """ + Tests that deleting a document automatically deletes its vector metadata due to cascading. + """ + # Create a document and a linked vector metadata entry + new_document = Document(user_id="cascade-test", title="Cascade Doc", text="Content") + db_session.add(new_document) + db_session.commit() + db_session.refresh(new_document) + + vector_meta = VectorMetadata(document_id=new_document.id, faiss_index=999, embedding_model="test") + db_session.add(vector_meta) + db_session.commit() + db_session.refresh(vector_meta) + + # Check that the vector metadata exists + assert db_session.query(VectorMetadata).filter(VectorMetadata.document_id == new_document.id).count() == 1 + + # Delete the document + db_session.delete(new_document) + db_session.commit() + + # Check that the document is gone and the vector metadata has been cascaded + assert db_session.query(Document).filter(Document.id == new_document.id).count() == 0 + assert db_session.query(VectorMetadata).filter(VectorMetadata.document_id == new_document.id).count() == 0 diff --git a/.gitignore b/.gitignore index ec8f767..c02da13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .env -**/.env \ No newline at end of file +**/.env +**/*.egg-info \ No newline at end of file diff --git a/KickoffPlan.md b/KickoffPlan.md index 3dab864..0e212f3 100644 --- a/KickoffPlan.md +++ b/KickoffPlan.md @@ -89,9 +89,3 @@ * **Syncing Data**: Ensure that metadata in **PostgreSQL** is synchronized with vectors in **FAISS** for accurate and consistent retrieval. - ---- - -This update to the plan leverages **PostgreSQL** for metadata management while **FAISS** handles efficient similarity search. The integration allows you to query and filter both metadata and vectors in an optimized manner, ensuring scalability and flexibility for future features. - -Let me know if you'd like more specifics or adjustments! diff --git a/README.md b/README.md index e69de29..b1f047a 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,106 @@ +Here is the **formatted and organized** version of your document for the **AI Model Hub Service**: + +--- + +# **AI Model Hub Service** + +This project is a central **hub service** designed to route requests to various **Large Language Models (LLMs)** and manage associated data. It integrates a **relational database (PostgreSQL)** for metadata with a **vector store (FAISS)** for efficient similarity search. + +--- + +## **High-Level Architecture** + +The service consists of several main components that work together to process user requests and manage data: + +* **API Server** + A **FastAPI** application that serves as the primary entry point for all client requests. + +* **LLM Router / Orchestrator** + The core logic that determines which LLM handles a request, manages chat history, and orchestrates interactions with the databases. + +* **Vector Database (FAISS + PostgreSQL)** + A two-part system for data storage: + + * **PostgreSQL**: Stores relational data like user sessions, chat messages, document metadata, and more. + * **FAISS**: An in-memory index for high-speed vector similarity search. + +--- + +## **Database Schema** + +The schema is designed to store **conversations**, **documents**, and the metadata necessary to link them together. + +### **1. sessions** + +Stores individual conversations with a user. + +* **Columns**: + `id`, `user_id`, `title`, `model_name`, `created_at`, `is_archived` + +--- + +### **2. messages** + +Stores each message within a session. + +* **Columns**: + `id`, `session_id`, `sender`, `content`, `created_at`, `model_response_time`, `token_count`, `metadata` + +--- + +### **3. documents** + +Stores the original text and metadata of a document for **RAG** (Retrieval-Augmented Generation). + +* **Columns**: + `id`, `title`, `text`, `source_url`, `author`, `status`, `created_at`, `user_id` + +--- + +### **4. vector\_metadata** + +Links a document to its vector representation in the **FAISS index**. + +* **Columns**: + `id`, `document_id`, `faiss_index`, `session_id`, `embedding_model` + +--- + +## **Technology Stack** + +* **API Framework**: FastAPI (Python) +* **LLM Integration**: LangChain (or similar) +* **Relational Database**: PostgreSQL +* **Vector Store**: FAISS +* **ORM**: SQLAlchemy +* **Testing Framework**: Pytest + +--- + +## **Running Tests** + +To ensure the database models work correctly, a set of tests is provided using **an in-memory SQLite database**. + +### **Steps to Run Tests** + +1. **Navigate to the project root.** + +2. **Set up your environment and install dependencies:** + + ```bash + pip install -e .[test] + ``` + + This installs the project in *editable mode* along with the test dependencies defined in `pyproject.toml`. + +3. **Run tests using Pytest:** + + ```bash + pytest + ``` + + Pytest will automatically discover and run all tests in the `tests/` directory. + +--- + +Let me know if you want this converted to **Markdown**, **PDF**, or **documentation site format (e.g., Sphinx or MkDocs)**. diff --git a/ai-hub/app/core/__init__.py b/ai-hub/app/core/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/core/__init__.py diff --git a/ai-hub/app/db/README.md b/ai-hub/app/db/README.md new file mode 100644 index 0000000..b77a001 --- /dev/null +++ b/ai-hub/app/db/README.md @@ -0,0 +1,120 @@ +Here is the properly formatted and organized version of your **AI Model Hub Database Schema** document: + +--- + +# **Database Schema for the AI Model Hub** + +This document provides a detailed overview of the **PostgreSQL database schema** for the AI Model Hub service, based on the SQLAlchemy models defined in `app/db/model.py`. The schema supports core functionalities like: + +* Conversational history +* Document storage +* Tracking RAG-related metadata + +--- + +## **Schema Diagram** + +The following diagram illustrates the relationships between the four main tables: `sessions`, `messages`, `documents`, and `vector_metadata`. + +``` ++-------------+ +-------------+ +| sessions | <------ | messages | ++-------------+ +-------------+ +| id | 1 | id | +| user_id | | session_id | +| title | | sender | +| model_name | | content | +| created_at | | created_at | +| is_archived | | model_response_time | ++-------------+ | token_count | + | metadata | + +-------------+ + ^ + | 1 ++------------------+ +-------------+ +| vector_metadata | ---1:1--| documents | ++------------------+ +-------------+ +| id | | id | +| document_id | | title | +| faiss_index | | text | +| session_id | | source_url | +| embedding_model | | author | ++------------------+ | status | + | created_at | + | user_id | + +-------------+ +``` + +--- + +## **Table Descriptions** + +### 1. **`sessions` Table** + +Stores metadata for each conversation session. Each row represents a single chat conversation. + +* `id` (Integer, Primary Key): Unique session identifier +* `user_id` (String): ID of the session's owner. Indexed for quick lookups +* `title` (String): Human-readable title for the session +* `model_name` (String): LLM used (e.g., `'gemini-1.5-pro'`, `'deepseek-chat'`) +* `created_at` (DateTime): Timestamp of session creation +* `is_archived` (Boolean): Soft delete flag + +--- + +### 2. **`messages` Table** + +Stores individual messages within a session. + +* `id` (Integer, Primary Key): Unique message identifier +* `session_id` (Integer, Foreign Key): Links to `sessions.id` +* `sender` (String): Either `'user'` or `'assistant'` +* `content` (Text): Message text +* `created_at` (DateTime): Timestamp of message +* `model_response_time` (Integer): Time (in seconds) to generate response +* `token_count` (Integer): Tokens used for the message +* `metadata` (JSON): Flexible field for model/tool-specific data + +--- + +### 3. **`documents` Table** + +Stores original text and metadata of documents ingested into the system. + +* `id` (Integer, Primary Key): Unique document identifier +* `title` (String): Document title +* `text` (Text): Full content of the document +* `source_url` (String): URL or file path of origin +* `author` (String): Author of the document +* `status` (String): `'processing'`, `'ready'`, or `'failed'` +* `created_at` (DateTime): Timestamp of upload +* `user_id` (String): ID of the uploading user + +--- + +### 4. **`vector_metadata` Table** + +Links documents to their vector representations and session context for RAG. + +* `id` (Integer, Primary Key): Unique metadata ID +* `document_id` (Integer, Foreign Key): Links to `documents.id` +* `faiss_index` (Integer): Index in the FAISS vector store +* `session_id` (Integer, Foreign Key): Session where this vector was used +* `embedding_model` (String): Embedding model used (e.g., `'text-embedding-004'`) + +--- + +## **Key Relationships** + +* **One-to-Many**: `sessions → messages` + A session contains multiple messages; each message belongs to one session. + +* **One-to-One**: `documents → vector_metadata` + Each document has a single vector metadata record. + +* **Many-to-One**: `vector_metadata → sessions` + Multiple vector metadata entries can reference the same session if used for RAG. + +--- + +Let me know if you’d like this in Markdown, PDF, or any specific format. diff --git a/ai-hub/app/db/__init__.py b/ai-hub/app/db/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/app/db/__init__.py diff --git a/ai-hub/app/db/database.py b/ai-hub/app/db/database.py new file mode 100644 index 0000000..2fc97fb --- /dev/null +++ b/ai-hub/app/db/database.py @@ -0,0 +1,54 @@ +import os +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT + +# --- Configuration --- +# Determines the database mode. Can be "postgres" or "sqlite". +# Defaults to "postgres" if not set. +DB_MODE = os.getenv("DB_MODE", "postgres").lower() + +# Default database URLs +POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" +SQLITE_DEFAULT_URL = "sqlite:///./ai_hub.db" + +DATABASE_URL = "" +engine_args = {} + +# --- Database Initialization --- +if DB_MODE == "sqlite": + print("✅ Initializing with SQLite in-file database.") + DATABASE_URL = SQLITE_DEFAULT_URL + # SQLite requires a specific argument to allow access from multiple threads, + # which is common in web applications. + engine_args = {"connect_args": {"check_same_thread": False}} +else: # Default to postgres + # Use the provided DATABASE_URL or fall back to the default. + DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) + DB_MODE = "postgres" + print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") + # pool_pre_ping checks if a connection is still alive before using it from the pool. + engine_args = {"pool_pre_ping": True} + + +# Create the SQLAlchemy engine with the determined settings +engine = create_engine(DATABASE_URL, **engine_args) + +# SessionLocal is a factory for creating new database session objects +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + +# Base is a class that our database model classes will inherit from. +Base = declarative_base() + + +# --- Dependency for FastAPI --- +def get_db(): + """ + FastAPI dependency that provides a database session for a single API request. + It ensures the session is always closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() + diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py new file mode 100644 index 0000000..ec4f703 --- /dev/null +++ b/ai-hub/app/db/models.py @@ -0,0 +1,159 @@ +from datetime import datetime +from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON +from sqlalchemy.orm import relationship + +# The declarative_base class is already defined in database.py. +# We will import it from there to ensure all models use the same base. +from .database import Base + +# --- SQLAlchemy Models --- +# These classes define the structure of the database tables and how they relate. + +class Session(Base): + """ + SQLAlchemy model for the 'sessions' table. + + Each session represents a single conversation between a user and the AI. + It links a user to a series of messages. + """ + __tablename__ = 'sessions' + + # Primary key for the session. + id = Column(Integer, primary_key=True, index=True) + # The ID of the user who owns this session. + user_id = Column(String, index=True, nullable=False) + # A title for the conversation, which can be generated by the AI. + title = Column(String, index=True, nullable=True) + # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek"). + model_name = Column(String, nullable=True) + # Timestamp for when the session was created. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # Flag to indicate if the session has been archived or soft-deleted. + is_archived = Column(Boolean, default=False, nullable=False) + + # Defines a one-to-many relationship with the Message table. + # 'back_populates' tells SQLAlchemy that there's a corresponding relationship + # on the other side. 'cascade' ensures that when a session is deleted, + # all its associated messages are also deleted. + messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Message(Base): + """ + SQLAlchemy model for the 'messages' table. + + This table stores the individual chat messages within a session, + including who sent them (user or AI) and the content. + """ + __tablename__ = 'messages' + + # Primary key for the message. + id = Column(Integer, primary_key=True, index=True) + # The foreign key that links this message to its parent session. + # This is a critical link for reconstructing chat history. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) + # Identifies the sender of the message, e.g., 'user' or 'assistant'. + sender = Column(String, nullable=False) + # The actual text content of the message. + content = Column(Text, nullable=False) + # Timestamp for when the message was sent. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # The time taken for the model to generate the response, in seconds. + model_response_time = Column(Integer, nullable=True) + # The number of tokens in the message (both input and output). + token_count = Column(Integer, nullable=True) + # A JSON field to store unstructured metadata about the message, such as tool calls. + # This column has been renamed from 'metadata' to avoid a conflict. + message_metadata = Column(JSON, nullable=True) + + + # Relationship back to the parent Session. + # This allows us to access the parent Session object from a Message object. + session = relationship("Session", back_populates="messages") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class Document(Base): + """ + SQLAlchemy model for the 'documents' table. + + This table stores the metadata and original text content of a document. + The content is the data that will be chunked, embedded, and used for RAG. + """ + __tablename__ = 'documents' + + # Primary key for the document, uniquely identifying each entry. + id = Column(Integer, primary_key=True, index=True) + # The title of the document for easy human-readable reference. + title = Column(String, index=True, nullable=False) + # The actual text content of the document. Using Text for potentially long strings. + text = Column(Text, nullable=False) + # The original source URL or path of the document. + source_url = Column(String, nullable=True) + # A string to identify the author of the document. + author = Column(String, nullable=True) + # The current processing status of the document (e.g., 'ready', 'processing', 'failed'). + status = Column(String, default="processing", nullable=False) + # Timestamp for when the document was added to the database. + created_at = Column(DateTime, default=datetime.utcnow, nullable=False) + # A string to identify the user who added the document, useful for multi-user apps. + user_id = Column(String, index=True, nullable=True) + + # Defines a one-to-one relationship with the VectorMetadata table. + vector_metadata = relationship( + "VectorMetadata", + back_populates="document", + cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted. + uselist=False + ) + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" + + +class VectorMetadata(Base): + """ + SQLAlchemy model for the 'vector_metadata' table. + + This table links a document to its corresponding vector representation + in the FAISS index. It is critical for syncing data between the + relational database and the vector store. + """ + __tablename__ = 'vector_metadata' + + # Primary key for the metadata entry. + id = Column(Integer, primary_key=True, index=True) + # Foreign key that links this metadata entry back to its Document. + document_id = Column(Integer, ForeignKey('documents.id'), unique=True) + # The index number in the FAISS vector store where the vector for this document is stored. + faiss_index = Column(Integer, nullable=False, index=True) + # Foreign key to link this vector metadata to a specific session. + # This is crucial for retrieving relevant RAG context for a given conversation. + session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True) + # The name of the embedding model used to create the vector. + embedding_model = Column(String, nullable=False) + + # Defines a many-to-one relationship with the Document table. + document = relationship("Document", back_populates="vector_metadata") + # Defines a many-to-one relationship with the Session table. + session = relationship("Session") + + def __repr__(self): + """ + Provides a helpful string representation of the object for debugging. + """ + return f"" diff --git a/ai-hub/pyproject.toml b/ai-hub/pyproject.toml new file mode 100644 index 0000000..2071533 --- /dev/null +++ b/ai-hub/pyproject.toml @@ -0,0 +1,5 @@ +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = [ + "tests" +] \ No newline at end of file diff --git a/ai-hub/requirements.txt b/ai-hub/requirements.txt index 4077554..9426b1e 100644 --- a/ai-hub/requirements.txt +++ b/ai-hub/requirements.txt @@ -6,6 +6,8 @@ pytest requests anyio +sqlalchemy +psycopg2 pytest-asyncio pytest-tornasync pytest-trio \ No newline at end of file diff --git a/ai-hub/setup.py b/ai-hub/setup.py new file mode 100644 index 0000000..5154feb --- /dev/null +++ b/ai-hub/setup.py @@ -0,0 +1,28 @@ +from setuptools import setup, find_packages + +# Read the requirements from requirements.txt +with open("requirements.txt") as f: + requirements = f.read().splitlines() + +setup( + name="ai-hub", + version="0.1.0", + description="An AI Model Hub Service with PostgreSQL and FAISS integration.", + author="Jerry Xie", + author_email="axieyangb@google.com", + packages=find_packages(), + install_requires=requirements, + # This is a key part of the setup. It defines the 'entry point' + # for the application, so you can run it with `uvicorn run:app`. + entry_points={ + "console_scripts": [ + "ai-hub-server=app.main:app" + ], + }, + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + ], + python_requires='>=3.11', +) \ No newline at end of file diff --git a/ai-hub/tests/db/test_database.py b/ai-hub/tests/db/test_database.py new file mode 100644 index 0000000..4f8921a --- /dev/null +++ b/ai-hub/tests/db/test_database.py @@ -0,0 +1,91 @@ +import os +import pytest +import importlib +from sqlalchemy.orm import Session +from sqlalchemy.exc import ResourceClosedError +from sqlalchemy import text +from unittest.mock import patch + +def test_sqlite_mode_initialization(monkeypatch): + """ + Tests if the database initializes in SQLite mode correctly. + """ + # Arrange: Set environment variable for SQLite mode + monkeypatch.setenv("DB_MODE", "sqlite") + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check if the configuration is correct for SQLite + assert database.DB_MODE == "sqlite" + assert "sqlite:///./ai_hub.db" in database.DATABASE_URL + assert "connect_args" in database.engine_args + assert database.engine_args["connect_args"] == {"check_same_thread": False} + + # Cleanup the created SQLite file after test, if it exists + if os.path.exists("ai_hub.db"): + os.remove("ai_hub.db") + +def test_postgres_mode_initialization(monkeypatch): + """ + Tests if the database initializes in PostgreSQL mode with a custom URL. + """ + # Arrange: Set env vars for PostgreSQL mode and a specific URL + monkeypatch.setenv("DB_MODE", "postgres") + monkeypatch.setenv("DATABASE_URL", "postgresql://test_user:test_password@testhost/test_db") + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check if the configuration is correct for PostgreSQL + assert database.DB_MODE == "postgres" + assert database.DATABASE_URL == "postgresql://test_user:test_password@testhost/test_db" + assert "pool_pre_ping" in database.engine_args + +def test_default_to_postgres_mode(monkeypatch): + """ + Tests if the system defaults to PostgreSQL mode when DB_MODE is not set. + """ + # Arrange: Ensure DB_MODE is not set + monkeypatch.delenv("DB_MODE", raising=False) + + # Act: Reload the module to apply the monkeypatched env vars + from app.db import database + importlib.reload(database) + + # Assert: Check that it defaulted to postgres + assert database.DB_MODE == "postgres" + assert "postgresql://user:password@localhost/ai_hub_db" in database.DATABASE_URL + +@patch('app.db.database.SessionLocal') +def test_get_db_yields_and_closes_session(mock_session_local, monkeypatch): + """ + Tests if the get_db() dependency function yields a valid, active session + and correctly closes it afterward by mocking the session object. + """ + # Arrange: Get the actual get_db function from the module + from app.db import database + + # Configure the mock session returned by SessionLocal() + mock_session = mock_session_local.return_value + + db_generator = database.get_db() + + # Act + # 1. Get the session object from the generator + db_session_instance = next(db_generator) + + # Assert + # 2. Check that the yielded object is our mock session + assert db_session_instance is mock_session + mock_session.close.assert_not_called() # The session should not be closed yet + + # 3. Exhaust the generator to trigger the 'finally' block + with pytest.raises(StopIteration): + next(db_generator) + + # 4. Assert that the close() method was called exactly once. + mock_session.close.assert_called_once() + diff --git a/ai-hub/tests/db/test_models.py b/ai-hub/tests/db/test_models.py new file mode 100644 index 0000000..c2e928b --- /dev/null +++ b/ai-hub/tests/db/test_models.py @@ -0,0 +1,201 @@ +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm.exc import DetachedInstanceError +from app.db.models import Base, Session, Message, Document, VectorMetadata + +# Use an in-memory SQLite database for testing. This is fast and ensures +# each test starts with a clean slate. +SQLALCHEMY_DATABASE_URL = "sqlite:///:memory:" + +# Create a database engine +engine = create_engine( + SQLALCHEMY_DATABASE_URL, + connect_args={"check_same_thread": False} +) + +# Create a configured "Session" class +TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + + +@pytest.fixture(scope="function") +def db_session(): + """ + A pytest fixture that creates a new database session for a test. + It creates all tables at the start of the test and drops them at the end. + This ensures that each test is completely isolated. + """ + # Create all tables in the database + Base.metadata.create_all(bind=engine) + db = TestingSessionLocal() + try: + yield db + finally: + db.close() + # Drop all tables after the test is finished + Base.metadata.drop_all(bind=engine) + + +def test_create_all_tables_with_relationships(db_session): + """ + Tests if all tables are created successfully and can be inspected. + """ + # Verify that all tables from the Base metadata were created + assert "sessions" in Base.metadata.tables + assert "messages" in Base.metadata.tables + assert "documents" in Base.metadata.tables + assert "vector_metadata" in Base.metadata.tables + + # Check for foreign key constraints to ensure relationships are set up + session_table = Base.metadata.tables['sessions'] + message_table = Base.metadata.tables['messages'] + document_table = Base.metadata.tables['documents'] + vector_metadata_table = Base.metadata.tables['vector_metadata'] + + assert len(message_table.foreign_keys) == 1 + assert list(message_table.foreign_keys)[0].column.table.name == 'sessions' + + assert len(vector_metadata_table.foreign_keys) == 2 + fk_columns = [fk.column.table.name for fk in vector_metadata_table.foreign_keys] + assert 'documents' in fk_columns + assert 'sessions' in fk_columns + + +def test_create_and_retrieve_session(db_session): + """ + Tests the creation and retrieval of a Session object. + """ + # Create a new session object + new_session = Session(user_id="test-user-123", title="Test Session", model_name="gemini") + + # Add to session and commit to the database + db_session.add(new_session) + db_session.commit() + db_session.refresh(new_session) + + # Retrieve the session from the database by its ID + retrieved_session = db_session.query(Session).filter(Session.id == new_session.id).first() + + # Assert that the retrieved session matches the original + assert retrieved_session is not None + assert retrieved_session.user_id == "test-user-123" + assert retrieved_session.title == "Test Session" + assert retrieved_session.model_name == "gemini" + + +def test_create_message_with_session_relationship(db_session): + """ + Tests the creation of a Message and its relationship with a Session. + """ + # First, create a session + new_session = Session(user_id="test-user-123") + db_session.add(new_session) + db_session.commit() + db_session.refresh(new_session) + + # Then, create a message linked to that session + new_message = Message( + session_id=new_session.id, + sender="user", + content="Hello, world!", + token_count=3, + model_response_time=0 + ) + db_session.add(new_message) + db_session.commit() + db_session.refresh(new_message) + + # Retrieve the message and verify its relationship + retrieved_message = db_session.query(Message).filter(Message.id == new_message.id).first() + + assert retrieved_message is not None + assert retrieved_message.session_id == new_session.id + assert retrieved_message.content == "Hello, world!" + # Verify the relationship attribute works + assert retrieved_message.session.user_id == "test-user-123" + + +def test_create_document_and_vector_metadata_relationship(db_session): + """ + Tests the creation of a Document and its linked VectorMetadata. + """ + # Create a new document + new_document = Document( + user_id="test-user-123", + title="Sample Doc", + text="This is some sample text.", + status="ready" + ) + db_session.add(new_document) + db_session.commit() + db_session.refresh(new_document) + + # Create vector metadata linked to the document + vector_meta = VectorMetadata( + document_id=new_document.id, + faiss_index=123, + embedding_model="test-model" + ) + db_session.add(vector_meta) + db_session.commit() + db_session.refresh(vector_meta) + + # Retrieve the document and verify the relationship + retrieved_doc = db_session.query(Document).filter(Document.id == new_document.id).first() + assert retrieved_doc is not None + assert retrieved_doc.vector_metadata is not None + assert retrieved_doc.vector_metadata.faiss_index == 123 + + +def test_cascade_delete_session_and_messages(db_session): + """ + Tests that deleting a session automatically deletes its associated messages due to cascading. + """ + # Create a session and some messages + new_session = Session(user_id="cascade-test") + db_session.add(new_session) + db_session.commit() + db_session.refresh(new_session) + + message1 = Message(session_id=new_session.id, sender="user", content="Msg 1") + message2 = Message(session_id=new_session.id, sender="user", content="Msg 2") + db_session.add_all([message1, message2]) + db_session.commit() + + # Check that messages exist before deletion + assert db_session.query(Message).filter(Message.session_id == new_session.id).count() == 2 + + # Delete the session + db_session.delete(new_session) + db_session.commit() + + # Check that the session is gone and the messages have been cascaded + assert db_session.query(Session).filter(Session.id == new_session.id).count() == 0 + assert db_session.query(Message).filter(Message.session_id == new_session.id).count() == 0 + + +def test_cascade_delete_document_and_vector_metadata(db_session): + """ + Tests that deleting a document automatically deletes its vector metadata due to cascading. + """ + # Create a document and a linked vector metadata entry + new_document = Document(user_id="cascade-test", title="Cascade Doc", text="Content") + db_session.add(new_document) + db_session.commit() + db_session.refresh(new_document) + + vector_meta = VectorMetadata(document_id=new_document.id, faiss_index=999, embedding_model="test") + db_session.add(vector_meta) + db_session.commit() + db_session.refresh(vector_meta) + + # Check that the vector metadata exists + assert db_session.query(VectorMetadata).filter(VectorMetadata.document_id == new_document.id).count() == 1 + + # Delete the document + db_session.delete(new_document) + db_session.commit() + + # Check that the document is gone and the vector metadata has been cascaded + assert db_session.query(Document).filter(Document.id == new_document.id).count() == 0 + assert db_session.query(VectorMetadata).filter(VectorMetadata.document_id == new_document.id).count() == 0 diff --git a/ai-hub/tests/test_main.py b/ai-hub/tests/test_main.py index c3439a4..ee713b3 100644 --- a/ai-hub/tests/test_main.py +++ b/ai-hub/tests/test_main.py @@ -1,9 +1,9 @@ -# tests/test_main.py - from fastapi.testclient import TestClient -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, AsyncMock + +# Import the FastAPI app instance to create a test client from app.main import app -from app.llm_providers import get_llm_provider + # Create a TestClient instance based on our FastAPI app client = TestClient(app) @@ -13,34 +13,53 @@ assert response.status_code == 200 assert response.json() == {"status": "AI Model Hub is running!"} -@patch('app.main.client.chat.completions.create') -def test_chat_handler_success(mock_create): - """Test the /chat endpoint with a successful, mocked API call.""" - # Configure the mock to return a predictable response - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - mock_response.choices[0].message = MagicMock() - mock_response.choices[0].message.content = "This is a mock response from DeepSeek." - mock_create.return_value = mock_response - +@patch('app.main.get_llm_provider') +def test_chat_handler_success(mock_get_llm_provider): + """ + Test the /chat endpoint with a successful, mocked LLM response. + + We patch the get_llm_provider factory function to control the + behavior of the LLM provider instance it returns. + """ + # Configure a mock LLM provider instance with an async method + mock_provider = MagicMock() + mock_provider.generate_response = AsyncMock(return_value="This is a mock response from a provider.") + + # Configure our mocked factory function to return our mock provider + mock_get_llm_provider.return_value = mock_provider + # Make the request to our app response = client.post("/chat", json={"prompt": "Hello there"}) - + # Assert our app behaved as expected assert response.status_code == 200 - assert response.json() == {"response": "This is a mock response from DeepSeek."} - # Verify that the mocked function was called - mock_create.assert_called_once() + assert response.json()["response"] == "This is a mock response from a provider." + + # Verify that the mocked factory and its method were called + mock_get_llm_provider.assert_called_once_with("deepseek") + mock_provider.generate_response.assert_called_once_with("Hello there") -@patch('app.main.client.chat.completions.create') -def test_chat_handler_api_failure(mock_create): - """Test the /chat endpoint when the external API fails.""" - # Configure the mock to raise an exception - mock_create.side_effect = Exception("API connection error") - +@patch('app.main.get_llm_provider') +def test_chat_handler_api_failure(mock_get_llm_provider): + """ + Test the /chat endpoint when the external LLM API fails. + + We configure the mocked provider's generate_response method to raise an exception. + """ + # Configure a mock LLM provider instance with an async method that raises an exception + mock_provider = MagicMock() + mock_provider.generate_response = AsyncMock(side_effect=Exception("API connection error")) + + # Configure our mocked factory function to return our mock provider + mock_get_llm_provider.return_value = mock_provider + # Make the request to our app response = client.post("/chat", json={"prompt": "This request will fail"}) - + # Assert our app handles the error gracefully assert response.status_code == 500 - assert response.json() == {"detail": "Failed to get response from the model."} \ No newline at end of file + assert "An error occurred with the deepseek API" in response.json()["detail"] + + # Verify that the mocked factory and its method were called + mock_get_llm_provider.assert_called_once_with("deepseek") + mock_provider.generate_response.assert_called_once_with("This request will fail")