clickhouserag 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- clickhouserag-0.1.0/LICENSE +21 -0
- clickhouserag-0.1.0/PKG-INFO +18 -0
- clickhouserag-0.1.0/README.md +1 -0
- clickhouserag-0.1.0/clickhouserag/config/__init__.py +1 -0
- clickhouserag-0.1.0/clickhouserag/config/base_config.py +1 -0
- clickhouserag-0.1.0/clickhouserag/config/clickhouse_config.py +1 -0
- clickhouserag-0.1.0/clickhouserag/core/__init__.py +1 -0
- clickhouserag-0.1.0/clickhouserag/data_access/__init__.py +1 -0
- clickhouserag-0.1.0/clickhouserag/data_access/abstract_client.py +38 -0
- clickhouserag-0.1.0/clickhouserag/data_access/abstract_table.py +56 -0
- clickhouserag-0.1.0/clickhouserag/data_access/clickhouse_client.py +85 -0
- clickhouserag-0.1.0/clickhouserag/data_access/clickhouse_table.py +54 -0
- clickhouserag-0.1.0/clickhouserag/rag/__init__.py +1 -0
- clickhouserag-0.1.0/clickhouserag/rag/base.py +33 -0
- clickhouserag-0.1.0/clickhouserag/rag/manager.py +202 -0
- clickhouserag-0.1.0/clickhouserag/utils/__init__.py +1 -0
- clickhouserag-0.1.0/clickhouserag/vectorizers/__init__.py +1 -0
- clickhouserag-0.1.0/clickhouserag/vectorizers/base.py +13 -0
- clickhouserag-0.1.0/clickhouserag/vectorizers/manager.py +34 -0
- clickhouserag-0.1.0/pyproject.toml +46 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2024 Leonid Chesnikov
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: clickhouserag
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A Python library for creating RAG with Clickhouse.
|
|
5
|
+
License: MIT
|
|
6
|
+
Author: Леонид Чесников
|
|
7
|
+
Author-email: leonid.chesnikov@gmail.com
|
|
8
|
+
Requires-Python: >=3.11,<4.0
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
13
|
+
Requires-Dist: clickhouse-connect (>=0.7.16,<0.8.0)
|
|
14
|
+
Requires-Dist: clickhouse-driver (>=0.2.8,<0.3.0)
|
|
15
|
+
Requires-Dist: numpy (<2.0.0)
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# RAG-clickhouse-template
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# RAG-clickhouse-template
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Configuration package for ClickhouseRAG."""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Base configuration module for ClickhouseRAG."""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Clickhouse configuration module for ClickhouseRAG."""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Core package for ClickhouseRAG."""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Data access package for ClickhouseRAG."""
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
"""Abstract client module for Clickhouse data access."""
|
|
2
|
+
|
|
3
|
+
from abc import ABC, abstractmethod
|
|
4
|
+
from typing import Any, Dict, List, Optional
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class ClickhouseClient(ABC):
|
|
8
|
+
"""Abstract class for Clickhouse client."""
|
|
9
|
+
|
|
10
|
+
@abstractmethod
|
|
11
|
+
def connect(self) -> None:
|
|
12
|
+
"""Connect to the Clickhouse database."""
|
|
13
|
+
pass
|
|
14
|
+
|
|
15
|
+
@abstractmethod
|
|
16
|
+
def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Any:
|
|
17
|
+
"""Execute a query in the Clickhouse database."""
|
|
18
|
+
pass
|
|
19
|
+
|
|
20
|
+
@abstractmethod
|
|
21
|
+
def fetch_one(self, query: str, params: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
|
|
22
|
+
"""Fetch one result from the Clickhouse database."""
|
|
23
|
+
pass
|
|
24
|
+
|
|
25
|
+
@abstractmethod
|
|
26
|
+
def fetch_all(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
|
|
27
|
+
"""Fetch all results from the Clickhouse database."""
|
|
28
|
+
pass
|
|
29
|
+
|
|
30
|
+
@abstractmethod
|
|
31
|
+
def fetch_column(self, query: str, column: int = 0, params: Optional[Dict[str, Any]] = None) -> List[Any]:
|
|
32
|
+
"""Fetch a specific column from the Clickhouse database."""
|
|
33
|
+
pass
|
|
34
|
+
|
|
35
|
+
@abstractmethod
|
|
36
|
+
def close(self) -> None:
|
|
37
|
+
"""Close the connection to the Clickhouse database."""
|
|
38
|
+
pass
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
"""Abstract table module for Clickhouse data access."""
|
|
2
|
+
|
|
3
|
+
from abc import ABC, abstractmethod
|
|
4
|
+
from typing import Any, Dict, List, Optional
|
|
5
|
+
|
|
6
|
+
from clickhouserag.data_access.abstract_client import ClickhouseClient
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class ClickhouseTable(ABC):
|
|
10
|
+
"""Abstract class for Clickhouse table management."""
|
|
11
|
+
|
|
12
|
+
def __init__(self, client: ClickhouseClient, table_name: str) -> None:
|
|
13
|
+
"""Initialize ClickhouseTable.
|
|
14
|
+
|
|
15
|
+
Args:
|
|
16
|
+
----
|
|
17
|
+
client (ClickhouseClient): The Clickhouse client.
|
|
18
|
+
table_name (str): The name of the table.
|
|
19
|
+
|
|
20
|
+
"""
|
|
21
|
+
self.client = client
|
|
22
|
+
self.table_name = table_name
|
|
23
|
+
|
|
24
|
+
@abstractmethod
|
|
25
|
+
def insert(self, values: List[Dict[str, Any]]) -> None:
|
|
26
|
+
"""Insert values into the table."""
|
|
27
|
+
pass
|
|
28
|
+
|
|
29
|
+
@abstractmethod
|
|
30
|
+
def update(self, values: Dict[str, Any], conditions: Dict[str, Any]) -> None:
|
|
31
|
+
"""Update values in the table based on conditions."""
|
|
32
|
+
pass
|
|
33
|
+
|
|
34
|
+
@abstractmethod
|
|
35
|
+
def delete(self, conditions: Dict[str, Any]) -> None:
|
|
36
|
+
"""Delete values from the table based on conditions."""
|
|
37
|
+
pass
|
|
38
|
+
|
|
39
|
+
@abstractmethod
|
|
40
|
+
def search(
|
|
41
|
+
self, query: str, params: Optional[Dict[str, Any]] = None
|
|
42
|
+
) -> List[Dict[str, Any]]:
|
|
43
|
+
"""Search the table based on a query."""
|
|
44
|
+
pass
|
|
45
|
+
|
|
46
|
+
@abstractmethod
|
|
47
|
+
def fetch_all(
|
|
48
|
+
self, params: Optional[Dict[str, Any]] = None
|
|
49
|
+
) -> List[Dict[str, Any]]:
|
|
50
|
+
"""Fetch all values from the table."""
|
|
51
|
+
pass
|
|
52
|
+
|
|
53
|
+
@abstractmethod
|
|
54
|
+
def reset_table(self) -> None:
|
|
55
|
+
"""Reset the table."""
|
|
56
|
+
pass
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
"""Clickhouse client module for Clickhouse data access."""
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
from typing import Any, Dict, List, Optional
|
|
5
|
+
|
|
6
|
+
from clickhouse_driver import Client
|
|
7
|
+
|
|
8
|
+
from clickhouserag.data_access.abstract_client import ClickhouseClient
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class ClickhouseConnectClient(ClickhouseClient):
|
|
12
|
+
"""Clickhouse client implementation using clickhouse-driver."""
|
|
13
|
+
|
|
14
|
+
def __init__(self, host: str, port: int, username: str, password: str, database: str):
|
|
15
|
+
"""Initialize ClickhouseConnectClient.
|
|
16
|
+
|
|
17
|
+
Args:
|
|
18
|
+
----
|
|
19
|
+
host (str): The host of the Clickhouse server.
|
|
20
|
+
port (int): The port of the Clickhouse server.
|
|
21
|
+
username (str): The username for Clickhouse authentication.
|
|
22
|
+
password (str): The password for Clickhouse authentication.
|
|
23
|
+
database (str): The database to connect to.
|
|
24
|
+
|
|
25
|
+
"""
|
|
26
|
+
self.host = host
|
|
27
|
+
self.port = port
|
|
28
|
+
self.username = username
|
|
29
|
+
self.password = password
|
|
30
|
+
self.database = database
|
|
31
|
+
self.client = None
|
|
32
|
+
self.logger = logging.getLogger(__name__)
|
|
33
|
+
|
|
34
|
+
def connect(self) -> None:
|
|
35
|
+
"""Connect to the Clickhouse database."""
|
|
36
|
+
try:
|
|
37
|
+
self.client = Client(host=self.host, port=self.port, user=self.username, password=self.password, database=self.database)
|
|
38
|
+
except Exception as e:
|
|
39
|
+
self.logger.error(f"Failed to connect to Clickhouse: {e}")
|
|
40
|
+
raise
|
|
41
|
+
|
|
42
|
+
def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Any:
|
|
43
|
+
"""Execute a query in the Clickhouse database."""
|
|
44
|
+
try:
|
|
45
|
+
if params:
|
|
46
|
+
return self.client.execute(query, params)
|
|
47
|
+
return self.client.execute(query)
|
|
48
|
+
except Exception as e:
|
|
49
|
+
self.logger.error(f"Failed to execute query: {e}")
|
|
50
|
+
raise
|
|
51
|
+
|
|
52
|
+
def fetch_one(self, query: str, params: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
|
|
53
|
+
"""Fetch one result from the Clickhouse database."""
|
|
54
|
+
try:
|
|
55
|
+
result = self.execute_query(query, params)
|
|
56
|
+
return result[0] if result else None
|
|
57
|
+
except Exception as e:
|
|
58
|
+
self.logger.error(f"Failed to fetch one: {e}")
|
|
59
|
+
raise
|
|
60
|
+
|
|
61
|
+
def fetch_all(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
|
|
62
|
+
"""Fetch all results from the Clickhouse database."""
|
|
63
|
+
try:
|
|
64
|
+
return self.execute_query(query, params)
|
|
65
|
+
except Exception as e:
|
|
66
|
+
self.logger.error(f"Failed to fetch all: {e}")
|
|
67
|
+
raise
|
|
68
|
+
|
|
69
|
+
def fetch_column(self, query: str, column: int = 0, params: Optional[Dict[str, Any]] = None) -> List[Any]:
|
|
70
|
+
"""Fetch a specific column from the Clickhouse database."""
|
|
71
|
+
try:
|
|
72
|
+
result = self.execute_query(query, params)
|
|
73
|
+
return [row[column] for row in result]
|
|
74
|
+
except Exception as e:
|
|
75
|
+
self.logger.error(f"Failed to fetch column: {e}")
|
|
76
|
+
raise
|
|
77
|
+
|
|
78
|
+
def close(self) -> None:
|
|
79
|
+
"""Close the connection to the Clickhouse database."""
|
|
80
|
+
try:
|
|
81
|
+
if self.client:
|
|
82
|
+
self.client.disconnect()
|
|
83
|
+
except Exception as e:
|
|
84
|
+
self.logger.error(f"Failed to close connection: {e}")
|
|
85
|
+
raise
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
"""Clickhouse table management module."""
|
|
2
|
+
|
|
3
|
+
from typing import Any, Dict, List, Optional
|
|
4
|
+
|
|
5
|
+
from clickhouserag.data_access.abstract_table import ClickhouseTable
|
|
6
|
+
from clickhouserag.data_access.clickhouse_client import ClickhouseConnectClient
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class ClickhouseTableManager(ClickhouseTable):
|
|
10
|
+
"""Clickhouse table manager implementation."""
|
|
11
|
+
|
|
12
|
+
def __init__(self, client: ClickhouseConnectClient, table_name: str) -> None:
|
|
13
|
+
"""Initialize ClickhouseTableManager.
|
|
14
|
+
|
|
15
|
+
Args:
|
|
16
|
+
----
|
|
17
|
+
client (ClickhouseConnectClient): The Clickhouse client.
|
|
18
|
+
table_name (str): The name of the table.
|
|
19
|
+
|
|
20
|
+
"""
|
|
21
|
+
super().__init__(client, table_name)
|
|
22
|
+
|
|
23
|
+
def insert(self, values: List[Dict[str, Any]]) -> None:
|
|
24
|
+
"""Insert values into the table."""
|
|
25
|
+
query = f"INSERT INTO {self.table_name} VALUES"
|
|
26
|
+
self.client.execute_query(query, values)
|
|
27
|
+
|
|
28
|
+
def update(self, values: Dict[str, Any], conditions: Dict[str, Any]) -> None:
|
|
29
|
+
"""Update values in the table based on conditions."""
|
|
30
|
+
set_clause = ", ".join([f"{key} = %({key})s" for key in values.keys()])
|
|
31
|
+
condition_clause = " AND ".join([f"{key} = %({key})s" for key in conditions.keys()])
|
|
32
|
+
query = f"ALTER TABLE {self.table_name} UPDATE {set_clause} WHERE {condition_clause}"
|
|
33
|
+
self.client.execute_query(query, {**values, **conditions})
|
|
34
|
+
|
|
35
|
+
def delete(self, conditions: Dict[str, Any]) -> None:
|
|
36
|
+
"""Delete values from the table based on conditions."""
|
|
37
|
+
condition_clause = " AND ".join([f"{key} = %(cond_{key})s" for key in conditions.keys()])
|
|
38
|
+
query = f"DELETE FROM {self.table_name} WHERE {condition_clause}"
|
|
39
|
+
params = {f"cond_{key}": value for key, value in conditions.items()}
|
|
40
|
+
self.client.execute_query(query, params)
|
|
41
|
+
|
|
42
|
+
def search(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
|
|
43
|
+
"""Search the table based on a query."""
|
|
44
|
+
return self.client.fetch_all(query, params)
|
|
45
|
+
|
|
46
|
+
def fetch_all(self, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
|
|
47
|
+
"""Fetch all values from the table."""
|
|
48
|
+
query = f"SELECT * FROM {self.table_name}"
|
|
49
|
+
return self.client.fetch_all(query, params)
|
|
50
|
+
|
|
51
|
+
def reset_table(self) -> None:
|
|
52
|
+
"""Reset the table."""
|
|
53
|
+
query = f"TRUNCATE TABLE {self.table_name}"
|
|
54
|
+
self.client.execute_query(query)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""RAG package for ClickhouseRAG."""
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
"""Base module for RAG management."""
|
|
2
|
+
|
|
3
|
+
from abc import ABC, abstractmethod
|
|
4
|
+
from typing import Any, Dict, List
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class RAGBase(ABC):
|
|
8
|
+
"""Abstract base class for RAG management."""
|
|
9
|
+
|
|
10
|
+
@abstractmethod
|
|
11
|
+
def add_data(self, data: Dict[str, Any]) -> None:
|
|
12
|
+
"""Add data to the RAG."""
|
|
13
|
+
pass
|
|
14
|
+
|
|
15
|
+
@abstractmethod
|
|
16
|
+
def delete_data(self, data_id: str) -> None:
|
|
17
|
+
"""Delete data from the RAG."""
|
|
18
|
+
pass
|
|
19
|
+
|
|
20
|
+
@abstractmethod
|
|
21
|
+
def update_data(self, data_id: str, new_data: Dict[str, Any]) -> None:
|
|
22
|
+
"""Update data in the RAG."""
|
|
23
|
+
pass
|
|
24
|
+
|
|
25
|
+
@abstractmethod
|
|
26
|
+
def search(self, query: str) -> List[Dict[str, Any]]:
|
|
27
|
+
"""Search the RAG."""
|
|
28
|
+
pass
|
|
29
|
+
|
|
30
|
+
@abstractmethod
|
|
31
|
+
def reset_database(self) -> None:
|
|
32
|
+
"""Reset the RAG database."""
|
|
33
|
+
pass
|
|
@@ -0,0 +1,202 @@
|
|
|
1
|
+
"""RAG manager module for managing data in Clickhouse."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import logging
|
|
5
|
+
from typing import Any, Dict, List, Optional
|
|
6
|
+
|
|
7
|
+
import numpy as np
|
|
8
|
+
|
|
9
|
+
from clickhouserag.data_access.clickhouse_client import ClickhouseClient
|
|
10
|
+
from clickhouserag.data_access.clickhouse_table import ClickhouseTableManager
|
|
11
|
+
from clickhouserag.rag.base import RAGBase
|
|
12
|
+
from clickhouserag.vectorizers.base import VectorizerBase
|
|
13
|
+
from clickhouserag.vectorizers.manager import VectorizerManager
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class RAGManager(RAGBase):
|
|
17
|
+
"""Manager class for RAG operations in Clickhouse."""
|
|
18
|
+
|
|
19
|
+
def __init__(self, client: ClickhouseClient, table_name: str, table_schema: Optional[Dict[str, str]] = None, engine: str = "MergeTree", order_by: str = "id") -> None:
|
|
20
|
+
"""Initialize RAGManager.
|
|
21
|
+
|
|
22
|
+
Args:
|
|
23
|
+
----
|
|
24
|
+
client (ClickhouseClient): The Clickhouse client.
|
|
25
|
+
table_name (str): The name of the table.
|
|
26
|
+
table_schema (Optional[Dict[str, str]]): The schema of the table.
|
|
27
|
+
engine (str): The engine to use for the table.
|
|
28
|
+
order_by (str): The column to order the table by.
|
|
29
|
+
|
|
30
|
+
"""
|
|
31
|
+
self.client = client
|
|
32
|
+
self.table_manager = ClickhouseTableManager(client, table_name)
|
|
33
|
+
self.table_name = table_name
|
|
34
|
+
self.table_schema = table_schema
|
|
35
|
+
self.engine = engine
|
|
36
|
+
self.order_by = order_by
|
|
37
|
+
self.vectorizer_manager = VectorizerManager()
|
|
38
|
+
self.logger = logging.getLogger(__name__)
|
|
39
|
+
|
|
40
|
+
if table_schema:
|
|
41
|
+
self._initialize_table(table_schema, engine, order_by)
|
|
42
|
+
|
|
43
|
+
def _initialize_table(self, table_schema: Dict[str, str], engine: str, order_by: str) -> None:
|
|
44
|
+
"""Initialize the table in Clickhouse based on the provided schema if it does not exist."""
|
|
45
|
+
if not self._check_table_exists():
|
|
46
|
+
self._create_table(table_schema, engine, order_by)
|
|
47
|
+
else:
|
|
48
|
+
self.logger.info(f"Table '{self.table_name}' already exists.")
|
|
49
|
+
|
|
50
|
+
def _create_table(self, table_schema: Dict[str, str], engine: str, order_by: str) -> None:
|
|
51
|
+
"""Create a table in Clickhouse based on the provided schema."""
|
|
52
|
+
try:
|
|
53
|
+
fields = ", ".join([f"{name} {dtype}" for name, dtype in table_schema.items()])
|
|
54
|
+
query = f"CREATE TABLE {self.table_name} ({fields}) ENGINE = {engine} ORDER BY {order_by}"
|
|
55
|
+
self.client.execute_query(query)
|
|
56
|
+
self.logger.info(f"Table '{self.table_name}' created with schema: {table_schema}, engine: {engine}, order by: {order_by}")
|
|
57
|
+
except Exception as e:
|
|
58
|
+
self.logger.error(f"Failed to create table '{self.table_name}': {e}")
|
|
59
|
+
raise
|
|
60
|
+
|
|
61
|
+
def _check_table_exists(self) -> bool:
|
|
62
|
+
"""Check if the table exists in the database."""
|
|
63
|
+
query = f"EXISTS TABLE {self.table_name}"
|
|
64
|
+
result = self.client.execute_query(query)
|
|
65
|
+
return result[0][0] == 1
|
|
66
|
+
|
|
67
|
+
def add_data(self, data: Dict[str, Any], vectorizer_name: Optional[str] = None) -> None:
|
|
68
|
+
"""Add data to the RAG."""
|
|
69
|
+
if vectorizer_name:
|
|
70
|
+
data["vector"] = self.vectorize(data, vectorizer_name)
|
|
71
|
+
self.validate_data(data)
|
|
72
|
+
self.table_manager.insert([data])
|
|
73
|
+
self.logger.info(f"Data added with id {data.get('id')}")
|
|
74
|
+
|
|
75
|
+
def add_bulk_data(self, data_list: List[Dict[str, Any]], vectorizer_name: Optional[str] = None) -> None:
|
|
76
|
+
"""Add multiple data records to the RAG."""
|
|
77
|
+
if vectorizer_name:
|
|
78
|
+
for data in data_list:
|
|
79
|
+
data["vector"] = self.vectorize(data, vectorizer_name)
|
|
80
|
+
for data in data_list:
|
|
81
|
+
self.validate_data(data)
|
|
82
|
+
self.table_manager.insert(data_list)
|
|
83
|
+
self.logger.info(f"Bulk data added with {len(data_list)} records")
|
|
84
|
+
|
|
85
|
+
def delete_data(self, data_id: str) -> None:
|
|
86
|
+
"""Delete data from the RAG."""
|
|
87
|
+
self.table_manager.delete({"id": data_id})
|
|
88
|
+
self.logger.info(f"Data deleted with id {data_id}")
|
|
89
|
+
|
|
90
|
+
def update_data(self, data_id: str, new_data: Dict[str, Any], vectorizer_name: Optional[str] = None) -> None:
|
|
91
|
+
"""Update data in the RAG."""
|
|
92
|
+
if vectorizer_name:
|
|
93
|
+
new_data["vector"] = self.vectorize(new_data, vectorizer_name)
|
|
94
|
+
self.validate_data(new_data)
|
|
95
|
+
self.table_manager.update(new_data, {"id": data_id})
|
|
96
|
+
self.logger.info(f"Data updated with id {data_id}")
|
|
97
|
+
|
|
98
|
+
def search(self, query: str, similarity: bool = False, top_k: Optional[int] = None) -> List[Dict[str, Any]]:
|
|
99
|
+
"""Search the RAG."""
|
|
100
|
+
if similarity:
|
|
101
|
+
return self.similarity_search(query, top_k)
|
|
102
|
+
else:
|
|
103
|
+
results = self.table_manager.search(query)
|
|
104
|
+
self.logger.info(f"Search executed with query '{query}', found {len(results)} results")
|
|
105
|
+
return results
|
|
106
|
+
|
|
107
|
+
def similarity_search(self, embedding: np.array, top_k: Optional[int]) -> List[Dict[str, Any]]:
|
|
108
|
+
"""Search the RAG based on cosine similarity."""
|
|
109
|
+
query = f"""
|
|
110
|
+
WITH %(embedding)s as query_vector
|
|
111
|
+
SELECT id, title,
|
|
112
|
+
arraySum(x -> x * x, vector) * arraySum(x -> x * x, query_vector) != 0
|
|
113
|
+
? arraySum((x, y) -> x * y, vector, query_vector) / sqrt(arraySum(x -> x * x, vector) * arraySum(x -> x * x, query_vector))
|
|
114
|
+
: 0 AS cosine_distance
|
|
115
|
+
FROM {self.table_name}
|
|
116
|
+
WHERE length(query_vector) == length(vector)
|
|
117
|
+
ORDER BY cosine_distance DESC
|
|
118
|
+
LIMIT %(top_k)s
|
|
119
|
+
"""
|
|
120
|
+
|
|
121
|
+
params = {"embedding": embedding.tolist(), "top_k": top_k}
|
|
122
|
+
result = self.client.execute_query(query, params=params, settings={"max_query_size": "10000000000000"})
|
|
123
|
+
self.logger.info(f"Similarity search executed with embedding, top {top_k} results found")
|
|
124
|
+
return [{"id": row[0], "title": row[1], "cosine_distance": row[2]} for row in result]
|
|
125
|
+
|
|
126
|
+
def compute_cosine_similarity(self, vector1: List[float], vector2: List[float]) -> float:
|
|
127
|
+
"""Compute the cosine similarity between two vectors."""
|
|
128
|
+
dot_product = sum(p * q for p, q in zip(vector1, vector2))
|
|
129
|
+
magnitude1 = sum(p ** 2 for p in vector1) ** 0.5
|
|
130
|
+
magnitude2 = sum(q ** 2 for q in vector2) ** 0.5
|
|
131
|
+
if magnitude1 * magnitude2 == 0:
|
|
132
|
+
return 0.0
|
|
133
|
+
return dot_product / (magnitude1 * magnitude2)
|
|
134
|
+
|
|
135
|
+
def reset_database(self) -> None:
|
|
136
|
+
"""Reset the RAG database."""
|
|
137
|
+
self.table_manager.reset_table()
|
|
138
|
+
self.logger.info("RAG table reset")
|
|
139
|
+
|
|
140
|
+
def backup_database(self, path: str) -> None:
|
|
141
|
+
"""Backup the RAG database to a file."""
|
|
142
|
+
data = self.table_manager.fetch_all()
|
|
143
|
+
with open(path, "w") as file:
|
|
144
|
+
json.dump(data, file)
|
|
145
|
+
self.logger.info(f"Database backup created at {path}")
|
|
146
|
+
|
|
147
|
+
def restore_database(self, path: str, table_schema: Optional[Dict[str, str]] = None, engine: str = "MergeTree", order_by: str = "id") -> None:
|
|
148
|
+
"""Restore the RAG database from a file."""
|
|
149
|
+
if table_schema:
|
|
150
|
+
self._create_table(table_schema, engine, order_by)
|
|
151
|
+
with open(path, "r") as file:
|
|
152
|
+
data = json.load(file)
|
|
153
|
+
self.table_manager.reset_table()
|
|
154
|
+
self.table_manager.insert(data)
|
|
155
|
+
self.logger.info(f"Database restored from {path}")
|
|
156
|
+
|
|
157
|
+
def save_to_file(self, path: str) -> None:
|
|
158
|
+
"""Save the RAG database to a file."""
|
|
159
|
+
self.backup_database(path)
|
|
160
|
+
|
|
161
|
+
def load_from_file(self, path: str, table_schema: Optional[Dict[str, str]] = None, engine: str = "MergeTree", order_by: str = "id") -> None:
|
|
162
|
+
"""Load the RAG database from a file."""
|
|
163
|
+
self.restore_database(path, table_schema, engine, order_by)
|
|
164
|
+
|
|
165
|
+
def validate_data(self, data: Dict[str, Any]) -> None:
|
|
166
|
+
"""Validate the data to be inserted into the RAG."""
|
|
167
|
+
if "id" not in data:
|
|
168
|
+
raise ValueError("Data must contain an 'id' field")
|
|
169
|
+
# Additional data validation can be added here
|
|
170
|
+
|
|
171
|
+
def get_data(self, data_id: str) -> Optional[Dict[str, Any]]:
|
|
172
|
+
"""Get data from the RAG by ID."""
|
|
173
|
+
query = f"SELECT * FROM {self.table_name} WHERE id = %(data_id)s"
|
|
174
|
+
result = self.client.fetch_one(query, params={"data_id": data_id})
|
|
175
|
+
if result:
|
|
176
|
+
return dict(zip(self.table_manager.columns, result))
|
|
177
|
+
return None
|
|
178
|
+
|
|
179
|
+
def set_data(self, data: Dict[str, Any], vectorizer_name: Optional[str] = None) -> None:
|
|
180
|
+
"""Set data in the RAG."""
|
|
181
|
+
if "id" in data:
|
|
182
|
+
if self.get_data(data["id"]):
|
|
183
|
+
self.update_data(data["id"], data, vectorizer_name)
|
|
184
|
+
else:
|
|
185
|
+
self.add_data(data, vectorizer_name)
|
|
186
|
+
else:
|
|
187
|
+
raise ValueError("Data must contain an 'id' field")
|
|
188
|
+
|
|
189
|
+
def get_vectorizer(self, name: str) -> Optional[VectorizerBase]:
|
|
190
|
+
"""Get a vectorizer by name."""
|
|
191
|
+
return self.vectorizer_manager.get_vectorizer(name)
|
|
192
|
+
|
|
193
|
+
def add_vectorizer(self, name: str, vectorizer: VectorizerBase) -> None:
|
|
194
|
+
"""Add a vectorizer to the manager."""
|
|
195
|
+
self.vectorizer_manager.add_vectorizer(name, vectorizer)
|
|
196
|
+
|
|
197
|
+
def vectorize(self, data: Dict[str, Any], vectorizer_name: str) -> List[float]:
|
|
198
|
+
"""Vectorize data using the specified vectorizer."""
|
|
199
|
+
vectorizer = self.get_vectorizer(vectorizer_name)
|
|
200
|
+
if not vectorizer:
|
|
201
|
+
raise ValueError(f"Vectorizer '{vectorizer_name}' not found")
|
|
202
|
+
return vectorizer.vectorize(data["title"])
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Utilities package for ClickhouseRAG."""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Vectorizers package for ClickhouseRAG."""
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
"""Base module for vectorizers in ClickhouseRAG."""
|
|
2
|
+
|
|
3
|
+
from abc import ABC, abstractmethod
|
|
4
|
+
from typing import Any, List
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class VectorizerBase(ABC):
|
|
8
|
+
"""Abstract base class for vectorizers."""
|
|
9
|
+
|
|
10
|
+
@abstractmethod
|
|
11
|
+
def vectorize(self, data: Any) -> List[float]:
|
|
12
|
+
"""Convert data into a vector representation."""
|
|
13
|
+
pass
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
"""Manager module for vectorizers in ClickhouseRAG."""
|
|
2
|
+
|
|
3
|
+
from typing import Any, Dict, List, Optional
|
|
4
|
+
|
|
5
|
+
from clickhouserag.vectorizers.base import VectorizerBase
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class VectorizerManager:
|
|
9
|
+
"""Manager class for handling multiple vectorizers."""
|
|
10
|
+
|
|
11
|
+
def __init__(self):
|
|
12
|
+
"""Initialize VectorizerManager."""
|
|
13
|
+
self.vectorizers: Dict[str, VectorizerBase] = {}
|
|
14
|
+
|
|
15
|
+
def add_vectorizer(self, name: str, vectorizer: VectorizerBase) -> None:
|
|
16
|
+
"""Add a vectorizer to the manager."""
|
|
17
|
+
self.vectorizers[name] = vectorizer
|
|
18
|
+
|
|
19
|
+
def remove_vectorizer(self, name: str) -> None:
|
|
20
|
+
"""Remove a vectorizer from the manager."""
|
|
21
|
+
if name in self.vectorizers:
|
|
22
|
+
del self.vectorizers[name]
|
|
23
|
+
|
|
24
|
+
def get_vectorizer(self, name: str) -> Optional[VectorizerBase]:
|
|
25
|
+
"""Get a vectorizer by name."""
|
|
26
|
+
return self.vectorizers.get(name)
|
|
27
|
+
|
|
28
|
+
def vectorize(self, name: str, data: Any) -> List[float]:
|
|
29
|
+
"""Vectorize data using the specified vectorizer."""
|
|
30
|
+
vectorizer = self.get_vectorizer(name)
|
|
31
|
+
if vectorizer:
|
|
32
|
+
return vectorizer.vectorize(data)
|
|
33
|
+
else:
|
|
34
|
+
raise ValueError(f"Vectorizer '{name}' not found.")
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
[tool.poetry]
|
|
2
|
+
name = "clickhouserag"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "A Python library for creating RAG with Clickhouse."
|
|
5
|
+
authors = ["Леонид Чесников <leonid.chesnikov@gmail.com>"]
|
|
6
|
+
license = "MIT"
|
|
7
|
+
readme = "README.md"
|
|
8
|
+
|
|
9
|
+
[tool.poetry.dependencies]
|
|
10
|
+
python = "^3.11"
|
|
11
|
+
clickhouse-connect = "^0.7.16"
|
|
12
|
+
numpy = "<2.0.0"
|
|
13
|
+
clickhouse-driver = "^0.2.8"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
[tool.poetry.group.dev.dependencies]
|
|
17
|
+
ruff = "^0.5.4"
|
|
18
|
+
pytest = "^8.3.1"
|
|
19
|
+
transformers = "^4.42.4"
|
|
20
|
+
|
|
21
|
+
[build-system]
|
|
22
|
+
requires = ["poetry-core"]
|
|
23
|
+
build-backend = "poetry.core.masonry.api"
|
|
24
|
+
|
|
25
|
+
[tool.ruff]
|
|
26
|
+
lint.select = [
|
|
27
|
+
"A", # Pyflakes (checks for various errors)
|
|
28
|
+
"B", # pycodestyle (PEP 8)
|
|
29
|
+
"E", # pycodestyle (PEP 8)
|
|
30
|
+
"F", # Pyflakes (checks for various errors)
|
|
31
|
+
"W", # pycodestyle (PEP 8)
|
|
32
|
+
"I", # isort (import sorting)
|
|
33
|
+
"C", # McCabe (complexity)
|
|
34
|
+
"D", # pydocstyle (docstring style)
|
|
35
|
+
"N", # pep8-naming (naming conventions)
|
|
36
|
+
"PT", # flake8-pytest-style (pytest style)
|
|
37
|
+
"Q", # flake8-quotes (quote consistency)
|
|
38
|
+
"S", # flake8-bandit (security issues)
|
|
39
|
+
"T", # flake8-debugger (debug statements)
|
|
40
|
+
]
|
|
41
|
+
lint.ignore = ["E501"]
|
|
42
|
+
fix = true
|
|
43
|
+
|
|
44
|
+
[tool.pylint.messages_control]
|
|
45
|
+
disable = ["one-blank-line-before-class", "multi-line-summary-first-line"]
|
|
46
|
+
enable = ["multi-line-summary-second-line"]
|