From 3dc33817ba126668223db679998a7bbf69c56456 Mon Sep 17 00:00:00 2001
From: Vasilije <8619304+Vasilije1990@users.noreply.github.com>
Date: Sun, 29 Oct 2023 19:11:20 +0100
Subject: [PATCH] Added docs functionality
---
README.md | 131 +++++++++++++++-----
level_3/.env.template | 8 +-
level_3/Dockerfile | 1 +
level_3/api.py | 98 +++++++++++++--
level_3/create_database.py | 80 ++++++++++++
level_3/database/database.py | 14 ++-
level_3/docker-compose.yml | 11 ++
level_3/entrypoint.sh | 15 ++-
level_3/models/docs.py | 3 +-
level_3/models/operation.py | 2 +
level_3/models/testoutput.py | 1 +
level_3/poetry.lock | 8 +-
level_3/rag_test_manager.py | 113 +++++++++--------
level_3/scripts/create_database.py | 1 +
level_3/vectordb/basevectordb.py | 6 +
level_3/vectordb/loaders/loaders.py | 72 +++++++----
level_3/vectordb/vectordb.py | 26 +++-
level_3/wait-for-it.sh | 182 ++++++++++++++++++++++++++++
18 files changed, 652 insertions(+), 120 deletions(-)
create mode 100644 level_3/create_database.py
create mode 100644 level_3/wait-for-it.sh
diff --git a/README.md b/README.md
index 2af3e4704..b7b1caaa4 100644
--- a/README.md
+++ b/README.md
@@ -1,20 +1,20 @@
# PromethAI-Memory
-Memory management and testing for the AI Applications and RAGs
-Dynamic Graph Memory Manager + DB + Rag Test Manager
+AI Applications and RAGs - Cognitive Architecture, Testability, Production Ready Apps
-
+
+
- 
+
-Open-source framework that manages memory for AI Agents and LLM apps
+Open-source framework for building and testing RAGs and Cognitive Architectures, designed for accuracy, transparency, and control.
-
+
@@ -52,9 +52,9 @@ Dynamic Graph Memory Manager + DB + Rag Test Manager
[//]: # (
)
-Share promethAI Repository
+Share promethAI Repository
-
+
@@ -71,33 +71,40 @@ Dynamic Graph Memory Manager + DB + Rag Test Manager
-
+
+This repo is built to test and evolve RAG architecture, inspired by human cognitive processes, using Python. It's aims to be production ready, testable, but give great visibility in how we build RAG applications.
+
+This project is a part of the [PromethAI](https://prometh.ai/) ecosystem.
+
+It runs in iterations, with each iteration building on the previous one.
+
+_Keep Ithaka always in your mind.
+Arriving there is what you’re destined for.
+But don’t hurry the journey at all.
+Better if it lasts for years_
-## Production-ready modern data platform
+### Installation
+To get started with PromethAI Memory, start with the latest iteration, and follow the instructions in the README.md file
-Browsing the database of theresanaiforthat.com, we can observe around [7000 new, mostly semi-finished projects](https://theresanaiforthat.com/) in the field of applied AI.
-It seems it has never been easier to create a startup, build an app, and go to market… and fail.
+### Current Focus
-Decades of technological advancements have led to small teams being able to do in 2023 what in 2015 required a team of dozens.
-Yet, the AI apps currently being pushed out still mostly feel and perform like demos.
-The rise of this new profession is perhaps signaling the need for a solution that is not yet there — a solution that in its essence represents a Large Language Model (LLM) — [a powerful general problem solver](https://lilianweng.github.io/posts/2023-06-23-agent/?fbclid=IwAR1p0W-Mg_4WtjOCeE8E6s7pJZlTDCDLmcXqHYVIrEVisz_D_S8LfN6Vv20) — available in the palm of your hand 24/7/365.
+RAG test manager can be used via API or via the CLI
-To address this issue, [dlthub](https://dlthub.com/) and [prometh.ai](http://prometh.ai/) will collaborate on a productionizing a common use-case, progressing step by step. We will utilize the LLMs, frameworks, and services, refining the code until we attain a clearer understanding of what a modern LLM architecture stack might entail.
+
-## Read more on our blog post [prometh.ai](http://prometh.ai/promethai-memory-blog-post-on)
+### Project Structure
-
-## Project Structure
-
-### Level 1 - OpenAI functions + Pydantic + DLTHub
+#### Level 1 - OpenAI functions + Pydantic + DLTHub
Scope: Give PDFs to the model and get the output in a structured format
+Blog post: https://prometh.ai/promethai-memory-blog-post-one
We introduce the following concepts:
- Structured output with Pydantic
- CMD script to process custom PDFs
-### Level 2 - Memory Manager + Metadata management
+#### Level 2 - Memory Manager + Metadata management
Scope: Give PDFs to the model and consolidate with the previous user activity and more
+Blog post: https://prometh.ai/promethai-memory-blog-post-two
We introduce the following concepts:
- Long Term Memory -> store and format the data
@@ -106,8 +113,9 @@ We introduce the following concepts:
- Docker
- API
-### Level 3 - Dynamic Graph Memory Manager + DB + Rag Test Manager
+#### Level 3 - Dynamic Graph Memory Manager + DB + Rag Test Manager
Scope: Store the data in N-related stores and test the retrieval with the Rag Test Manager
+Blog post: https://prometh.ai/promethai-memory-blog-post-three
- Dynamic Memory Manager -> store the data in N hierarchical stores
- Auto-generation of tests
- Multiple file formats supported
@@ -116,26 +124,92 @@ Scope: Store the data in N-related stores and test the retrieval with the Rag Te
- API
-## Run the level 3
+### Run the level 3
Make sure you have Docker, Poetry, and Python 3.11 installed and postgres installed.
-Copy the .env.example to .env and fill the variables
+Copy the .env.example to .env and fill in the variables
-Start the docker:
+Two ways to run the level 3:
-```docker compose up promethai_mem ```
+#### Docker:
+
+Copy the .env.template to .env and fill in the variables
+Specify the environment variable in the .env file to "docker"
+
+
+Launch the docker image:
+
+```docker compose up promethai_mem ```
+
+Send the request to the API:
+
+```
+curl -X POST -H "Content-Type: application/json" -d '{
+ "payload": {
+ "user_id": "681",
+ "data": [".data/3ZCCCW.pdf"],
+ "test_set": "sample",
+ "params": ["chunk_size"],
+ "metadata": "sample",
+ "retriever_type": "single_document_context"
+ }
+}' http://0.0.0.0:8000/rag-test/rag_test_run
+
+```
+Params:
+
+data -> list of URLs or path to the file, located in the .data folder (pdf, docx, txt, html)
+test_set -> sample, manual (list of questions and answers)
+metadata -> sample, manual (json) or version (in progress)
+params -> chunk_size, chunk_overlap, search_type (hybrid, bm25), embeddings
+retriever_type -> llm_context, single_document_context, multi_document_context, cognitive_architecture(coming soon)
+
+Inspect the results in the DB:
+
+``` docker exec -it postgres psql -U bla ```
+
+``` \c bubu ```
+
+``` select * from test_outputs; ```
+
+Or set up the superset to visualize the results:
+
+
+
+#### Poetry environment:
+
+
+Copy the .env.template to .env and fill in the variables
+Specify the environment variable in the .env file to "local"
Use the poetry environment:
``` poetry shell ```
+Change the .env file Environment variable to "local"
+
+Launch the postgres DB
+
+``` docker compose up postgres ```
+
+Launch the superset
+
+``` docker compose up superset ```
+
+Open the superset in your browser
+
+``` http://localhost:8088 ```
+Add the Postgres datasource to the Superset with the following connection string:
+
+``` postgres://bla:bla@postgres:5432/bubu ```
+
Make sure to run to initialize DB tables
``` python scripts/create_database.py ```
-After that, you can run the RAG test manager.
+After that, you can run the RAG test manager from your command line.
```
@@ -149,3 +223,4 @@ After that, you can run the RAG test manager.
```
Examples of metadata structure and test set are in the folder "example_data"
+
diff --git a/level_3/.env.template b/level_3/.env.template
index bc9893356..8ca7daed0 100644
--- a/level_3/.env.template
+++ b/level_3/.env.template
@@ -1,3 +1,9 @@
OPENAI_API_KEY=sk
WEAVIATE_URL =
-WEAVIATE_API_KEY =
\ No newline at end of file
+WEAVIATE_API_KEY =
+ENVIRONMENT = docker
+POSTGRES_USER = bla
+POSTGRES_PASSWORD = bla
+POSTGRES_DB = bubu
+POSTGRES_HOST = localhost
+POSTGRES_HOST_DOCKER = postgres
\ No newline at end of file
diff --git a/level_3/Dockerfile b/level_3/Dockerfile
index f2b0a3efa..feb677fe3 100644
--- a/level_3/Dockerfile
+++ b/level_3/Dockerfile
@@ -43,6 +43,7 @@ RUN apt-get update -q && \
WORKDIR /app
COPY . /app
+COPY scripts/ /app
COPY entrypoint.sh /app/entrypoint.sh
COPY scripts/create_database.py /app/create_database.py
RUN chmod +x /app/entrypoint.sh
diff --git a/level_3/api.py b/level_3/api.py
index 953150d61..aa200d56a 100644
--- a/level_3/api.py
+++ b/level_3/api.py
@@ -1,9 +1,11 @@
+import json
import logging
import os
+from enum import Enum
from typing import Dict, Any
import uvicorn
-from fastapi import FastAPI
+from fastapi import FastAPI, BackgroundTasks
from fastapi.responses import JSONResponse
from pydantic import BaseModel
@@ -11,6 +13,7 @@ from database.database import AsyncSessionLocal
from database.database_crud import session_scope
from vectorstore_manager import Memory
from dotenv import load_dotenv
+from rag_test_manager import start_test
# Set up logging
logging.basicConfig(
@@ -200,25 +203,100 @@ def memory_factory(memory_type):
memory_list = ["episodic", "buffer", "semantic"]
for memory_type in memory_list:
memory_factory(memory_type)
+class TestSetType(Enum):
+ SAMPLE = "sample"
+ MANUAL = "manual"
+def get_test_set(test_set_type, folder_path="example_data", payload=None):
+ if test_set_type == TestSetType.SAMPLE:
+ file_path = os.path.join(folder_path, "test_set.json")
+ if os.path.isfile(file_path):
+ with open(file_path, "r") as file:
+ return json.load(file)
+ elif test_set_type == TestSetType.MANUAL:
+ # Check if the manual test set is provided in the payload
+ if payload and "manual_test_set" in payload:
+ return payload["manual_test_set"]
+ else:
+ # Attempt to load the manual test set from a file
+ pass
+
+ return None
+
+
+class MetadataType(Enum):
+ SAMPLE = "sample"
+ MANUAL = "manual"
+
+def get_metadata(metadata_type, folder_path="example_data", payload=None):
+ if metadata_type == MetadataType.SAMPLE:
+ file_path = os.path.join(folder_path, "metadata.json")
+ if os.path.isfile(file_path):
+ with open(file_path, "r") as file:
+ return json.load(file)
+ elif metadata_type == MetadataType.MANUAL:
+ # Check if the manual metadata is provided in the payload
+ if payload and "manual_metadata" in payload:
+ return payload["manual_metadata"]
+ else:
+ pass
+
+ return None
@app.post("/rag-test/rag_test_run", response_model=dict)
async def rag_test_run(
payload: Payload,
- # files: List[UploadFile] = File(...),
+ background_tasks: BackgroundTasks,
):
try:
- from rag_test_manager import start_test
- logging.info(" Running RAG Test ")
+ logging.info("Starting RAG Test")
decoded_payload = payload.payload
- output = await start_test(data=decoded_payload['data'], test_set=decoded_payload['test_set'], user_id=decoded_payload['user_id'], params=decoded_payload['params'], metadata=decoded_payload['metadata'],
- retriever_type=decoded_payload['retriever_type'])
- return JSONResponse(content={"response": output}, status_code=200)
- except Exception as e:
- return JSONResponse(
- content={"response": {"error": str(e)}}, status_code=503
+ test_set_type = TestSetType(decoded_payload['test_set'])
+
+ metadata_type = MetadataType(decoded_payload['metadata'])
+
+ metadata = get_metadata(metadata_type, payload=decoded_payload)
+ if metadata is None:
+ return JSONResponse(content={"response": "Invalid metadata value"}, status_code=400)
+
+ test_set = get_test_set(test_set_type, payload=decoded_payload)
+ if test_set is None:
+ return JSONResponse(content={"response": "Invalid test_set value"}, status_code=400)
+
+ async def run_start_test(data, test_set, user_id, params, metadata, retriever_type):
+ result = await start_test(data = data, test_set = test_set, user_id =user_id, params =params, metadata =metadata, retriever_type=retriever_type)
+
+ logging.info("Retriever DATA type", type(decoded_payload['data']))
+
+ background_tasks.add_task(
+ run_start_test,
+ decoded_payload['data'],
+ test_set,
+ decoded_payload['user_id'],
+ decoded_payload['params'],
+ metadata,
+ decoded_payload['retriever_type']
)
+ logging.info("Retriever type", decoded_payload['retriever_type'])
+ return JSONResponse(content={"response": "Task has been started"}, status_code=200)
+
+ except Exception as e:
+ return JSONResponse(
+
+ content={"response": {"error": str(e)}}, status_code=503
+
+ )
+
+
+# @app.get("/rag-test/{task_id}")
+# async def check_task_status(task_id: int):
+# task_status = task_status_db.get(task_id, "not_found")
+#
+# if task_status == "not_found":
+# return {"status": "Task not found"}
+#
+# return {"status": task_status}
# @app.get("/available-buffer-actions", response_model=dict)
# async def available_buffer_actions(
diff --git a/level_3/create_database.py b/level_3/create_database.py
new file mode 100644
index 000000000..3d8ff426d
--- /dev/null
+++ b/level_3/create_database.py
@@ -0,0 +1,80 @@
+import sys
+import os
+
+# this is needed to import classes from other modules
+script_dir = os.path.dirname(os.path.abspath(__file__))
+# Get the parent directory of your script and add it to sys.path
+parent_dir = os.path.dirname(script_dir)
+sys.path.append(parent_dir)
+
+from database.database import Base, engine
+import models.memory
+import models.metadatas
+import models.operation
+import models.sessions
+import models.testoutput
+import models.testset
+import models.user
+import models.docs
+from sqlalchemy import create_engine, text
+import psycopg2
+from dotenv import load_dotenv
+load_dotenv()
+import os
+
+
+
+
+
+
+def create_admin_engine(username, password, host, database_name):
+ admin_url = f"postgresql://{username}:{password}@{host}:5432/{database_name}"
+ return create_engine(admin_url)
+
+
+def database_exists(username, password, host, db_name):
+ engine = create_admin_engine(username, password, host, db_name)
+ connection = engine.connect()
+ query = text(f"SELECT 1 FROM pg_database WHERE datname='{db_name}'")
+ result = connection.execute(query).fetchone()
+ connection.close()
+ engine.dispose()
+ return result is not None
+
+
+def create_database(username, password, host, db_name):
+ engine = create_admin_engine(username, password, host, db_name)
+ connection = engine.raw_connection()
+ connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+ cursor = connection.cursor()
+ cursor.execute(f"CREATE DATABASE {db_name}")
+ cursor.close()
+ connection.close()
+ engine.dispose()
+
+
+def create_tables(engine):
+ Base.metadata.create_all(bind=engine)
+
+if __name__ == "__main__":
+ username = os.getenv('POSTGRES_USER')
+ password = os.getenv('POSTGRES_PASSWORD')
+ database_name = os.getenv('POSTGRES_DB')
+ environment = os.environ.get("ENVIRONMENT")
+
+ if environment == "local":
+ host = os.getenv('POSTGRES_HOST')
+
+ elif environment == "docker":
+ host = os.getenv('POSTGRES_HOST_DOCKER')
+ else:
+ host = os.getenv('POSTGRES_HOST_DOCKER')
+
+ engine = create_admin_engine(username, password, host, database_name)
+
+ if not database_exists(username, password, host, database_name):
+ print(f"Database {database_name} does not exist. Creating...")
+ create_database(username, password, host, database_name)
+ print(f"Database {database_name} created successfully.")
+
+ create_tables(engine)
\ No newline at end of file
diff --git a/level_3/database/database.py b/level_3/database/database.py
index 122efb64e..b100f5c8c 100644
--- a/level_3/database/database.py
+++ b/level_3/database/database.py
@@ -24,7 +24,19 @@ RETRY_DELAY = 5
username = os.getenv('POSTGRES_USER')
password = os.getenv('POSTGRES_PASSWORD')
database_name = os.getenv('POSTGRES_DB')
-host = os.getenv('POSTGRES_HOST')
+import os
+
+environment = os.environ.get("ENVIRONMENT")
+
+if environment == "local":
+ host= os.getenv('POSTGRES_HOST')
+
+elif environment == "docker":
+ host= os.getenv('POSTGRES_HOST_DOCKER')
+else:
+ host= os.getenv('POSTGRES_HOST_DOCKER')
+
+
# Use the asyncpg driver for async operation
SQLALCHEMY_DATABASE_URL = f"postgresql+asyncpg://{username}:{password}@{host}:5432/{database_name}"
diff --git a/level_3/docker-compose.yml b/level_3/docker-compose.yml
index 8e7e56698..97564a17e 100644
--- a/level_3/docker-compose.yml
+++ b/level_3/docker-compose.yml
@@ -18,14 +18,25 @@ services:
- promethai_mem_backend
build:
context: ./
+
volumes:
- "./:/app"
+ - ./.data:/app/.data
+
environment:
- HOST=0.0.0.0
profiles: ["exclude-from-up"]
ports:
- 8000:8000
- 443:443
+ depends_on:
+ - postgres
+ deploy:
+ resources:
+ limits:
+ cpus: "4.0"
+ memory: 8GB
+
postgres:
image: postgres
diff --git a/level_3/entrypoint.sh b/level_3/entrypoint.sh
index 5a87898d6..4cf551d01 100755
--- a/level_3/entrypoint.sh
+++ b/level_3/entrypoint.sh
@@ -1,7 +1,20 @@
#!/bin/bash
export ENVIRONMENT
+# Run Python scripts with error handling
+echo "Running fetch_secret.py"
python fetch_secret.py
+if [ $? -ne 0 ]; then
+ echo "Error: fetch_secret.py failed"
+ exit 1
+fi
+
+echo "Running create_database.py"
python create_database.py
+if [ $? -ne 0 ]; then
+ echo "Error: create_database.py failed"
+ exit 1
+fi
# Start Gunicorn
-gunicorn -w 2 -k uvicorn.workers.UvicornWorker -t 120 --bind=0.0.0.0:8000 --bind=0.0.0.0:443 --log-level debug api:app
\ No newline at end of file
+echo "Starting Gunicorn"
+gunicorn -w 3 -k uvicorn.workers.UvicornWorker -t 30000 --bind=0.0.0.0:8000 --bind=0.0.0.0:443 --log-level debug api:app
diff --git a/level_3/models/docs.py b/level_3/models/docs.py
index 95d98c485..38166956b 100644
--- a/level_3/models/docs.py
+++ b/level_3/models/docs.py
@@ -15,4 +15,5 @@ class DocsModel(Base):
created_at = Column(DateTime, default=datetime.utcnow)
updated_at = Column(DateTime, onupdate=datetime.utcnow)
- operation = relationship("Operation", back_populates="docs")
+
+ operations = relationship("Operation", back_populates="docs")
\ No newline at end of file
diff --git a/level_3/models/operation.py b/level_3/models/operation.py
index 595745b15..1d06657d9 100644
--- a/level_3/models/operation.py
+++ b/level_3/models/operation.py
@@ -14,6 +14,7 @@ class Operation(Base):
id = Column(String, primary_key=True)
user_id = Column(String, ForeignKey('users.id'), index=True) # Link to User
operation_type = Column(String, nullable=True)
+ operation_status = Column(String, nullable=True)
operation_params = Column(String, nullable=True)
number_of_files = Column(Integer, nullable=True)
test_set_id = Column(String, ForeignKey('test_sets.id'), index=True)
@@ -24,6 +25,7 @@ class Operation(Base):
# Relationships
user = relationship("User", back_populates="operations")
test_set = relationship("TestSet", back_populates="operations")
+ docs = relationship("DocsModel", back_populates="operations")
def __repr__(self):
return f""
diff --git a/level_3/models/testoutput.py b/level_3/models/testoutput.py
index ed83decbe..4731a6e46 100644
--- a/level_3/models/testoutput.py
+++ b/level_3/models/testoutput.py
@@ -36,6 +36,7 @@ class TestOutput(Base):
test_output = Column(String, nullable=True)
test_expected_output = Column(String, nullable=True)
test_context = Column(String, nullable=True)
+ number_of_memories = Column(String, nullable=True)
test_results = Column(JSON, nullable=True)
created_at = Column(DateTime, default=datetime.utcnow)
diff --git a/level_3/poetry.lock b/level_3/poetry.lock
index 5bda9ee7e..8751910e2 100644
--- a/level_3/poetry.lock
+++ b/level_3/poetry.lock
@@ -3869,18 +3869,18 @@ diagrams = ["jinja2", "railroad-diagrams"]
[[package]]
name = "pypdf"
-version = "3.16.4"
+version = "3.17.0"
description = "A pure-python PDF library capable of splitting, merging, cropping, and transforming PDF files"
optional = false
python-versions = ">=3.6"
files = [
- {file = "pypdf-3.16.4-py3-none-any.whl", hash = "sha256:a9b1eaf2db4c2edd93093470d33c3f353235c4a694f8a426a92a8ce77cea9eb7"},
- {file = "pypdf-3.16.4.tar.gz", hash = "sha256:01927771b562d4ba84939ef95b393f0179166da786c5db710d07f807c52f480d"},
+ {file = "pypdf-3.17.0-py3-none-any.whl", hash = "sha256:67f6bb7acd8fdbcf7e7a7d5319d12b8de100f5f94538d6e5647aaec3eb7c7dde"},
+ {file = "pypdf-3.17.0.tar.gz", hash = "sha256:9fab275fea57c9e5b2416035d13d867a459ebe36294a4c39a3d0bb45a7404bad"},
]
[package.extras]
crypto = ["PyCryptodome", "cryptography"]
-dev = ["black", "flit", "pip-tools", "pre-commit (<2.18.0)", "pytest-cov", "pytest-socket", "pytest-timeout", "wheel"]
+dev = ["black", "flit", "pip-tools", "pre-commit (<2.18.0)", "pytest-cov", "pytest-socket", "pytest-timeout", "pytest-xdist", "wheel"]
docs = ["myst_parser", "sphinx", "sphinx_rtd_theme"]
full = ["Pillow (>=8.0.0)", "PyCryptodome", "cryptography"]
image = ["Pillow (>=8.0.0)"]
diff --git a/level_3/rag_test_manager.py b/level_3/rag_test_manager.py
index 63ce1f388..df0c86bfd 100644
--- a/level_3/rag_test_manager.py
+++ b/level_3/rag_test_manager.py
@@ -30,6 +30,7 @@ from models.testset import TestSet
from models.testoutput import TestOutput
from models.metadatas import MetaDatas
from models.operation import Operation
+from models.docs import DocsModel
load_dotenv()
import ast
@@ -90,7 +91,10 @@ def get_document_names(doc_input):
- Folder path: get_document_names(".data")
- Single document file path: get_document_names(".data/example.pdf")
- Document name provided as a string: get_document_names("example.docx")
+
"""
+ if isinstance(doc_input, list):
+ return doc_input
if os.path.isdir(doc_input):
# doc_input is a folder
folder_path = doc_input
@@ -114,7 +118,17 @@ async def add_entity(session, entity):
s.add(entity) # No need to commit; session_scope takes care of it
return "Successfully added entity"
+async def update_entity(session, model, entity_id, new_value):
+ async with session_scope(session) as s:
+ # Retrieve the entity from the database
+ entity = await s.get(model, entity_id)
+ if entity:
+ # Update the relevant column and 'updated_at' will be automatically updated
+ entity.operation_status = new_value
+ return "Successfully updated entity"
+ else:
+ return "Entity not found"
async def retrieve_job_by_id(session, user_id, job_id):
try:
@@ -313,8 +327,8 @@ async def eval_test(
test_case = LLMTestCase(
input=query,
actual_output=result_output,
- expected_output=expected_output,
- context=context,
+ expected_output=[expected_output],
+ context=[context],
)
metric = OverallScoreMetric()
@@ -472,20 +486,20 @@ async def start_test(
if params is None:
data_format = data_format_route(
- data
+ data[0]
) # Assume data_format_route is predefined
logging.info("Data format is %s", data_format)
- data_location = data_location_route(data)
+ data_location = data_location_route(data[0])
logging.info(
"Data location is %s", data_location
) # Assume data_location_route is predefined
test_params = generate_param_variants(included_params=["chunk_size"])
if params:
data_format = data_format_route(
- data
+ data[0]
) # Assume data_format_route is predefined
logging.info("Data format is %s", data_format)
- data_location = data_location_route(data)
+ data_location = data_location_route(data[0])
logging.info(
"Data location is %s", data_location
)
@@ -508,6 +522,7 @@ async def start_test(
user_id=user_id,
operation_params=str(test_params),
number_of_files=count_files_in_data_folder(),
+ operation_status = "RUNNING",
operation_type=retriever_type,
test_set_id=test_set_id,
),
@@ -517,7 +532,7 @@ async def start_test(
await add_entity(
session,
- Docs(
+ DocsModel(
id=str(uuid.uuid4()),
operation_id=job_id,
doc_name = doc
@@ -586,11 +601,13 @@ async def start_test(
return retrieve_action["data"]["Get"][test_id][0]["text"]
async def run_eval(test_item, search_result):
+ logging.info("Initiated test set evaluation")
test_eval = await eval_test(
- query=test_item["question"],
- expected_output=test_item["answer"],
+ query=str(test_item["question"]),
+ expected_output=str(test_item["answer"]),
context=str(search_result),
)
+ logging.info("Successfully evaluated test set")
return test_eval
async def run_generate_test_set(test_id):
@@ -607,9 +624,6 @@ async def start_test(
return dynamic_test_manager(retrieve_action)
test_eval_pipeline = []
-
-
-
if retriever_type == "llm_context":
for test_qa in test_set:
context = ""
@@ -690,47 +704,49 @@ async def start_test(
),
)
+ await update_entity(session, Operation, job_id, "COMPLETED")
+
return results
async def main():
- metadata = {
- "version": "1.0",
- "agreement_id": "AG123456",
- "privacy_policy": "https://example.com/privacy",
- "terms_of_service": "https://example.com/terms",
- "format": "json",
- "schema_version": "1.1",
- "checksum": "a1b2c3d4e5f6",
- "owner": "John Doe",
- "license": "MIT",
- "validity_start": "2023-08-01",
- "validity_end": "2024-07-31",
- }
-
- test_set = [
- {
- "question": "Who is the main character in 'The Call of the Wild'?",
- "answer": "Buck",
- },
- {"question": "Who wrote 'The Call of the Wild'?", "answer": "Jack London"},
- {
- "question": "Where does Buck live at the start of the book?",
- "answer": "In the Santa Clara Valley, at Judge Miller’s place.",
- },
- {
- "question": "Why is Buck kidnapped?",
- "answer": "He is kidnapped to be sold as a sled dog in the Yukon during the Klondike Gold Rush.",
- },
- {
- "question": "How does Buck become the leader of the sled dog team?",
- "answer": "Buck becomes the leader after defeating the original leader, Spitz, in a fight.",
- },
- ]
+ # metadata = {
+ # "version": "1.0",
+ # "agreement_id": "AG123456",
+ # "privacy_policy": "https://example.com/privacy",
+ # "terms_of_service": "https://example.com/terms",
+ # "format": "json",
+ # "schema_version": "1.1",
+ # "checksum": "a1b2c3d4e5f6",
+ # "owner": "John Doe",
+ # "license": "MIT",
+ # "validity_start": "2023-08-01",
+ # "validity_end": "2024-07-31",
+ # }
+ #
+ # test_set = [
+ # {
+ # "question": "Who is the main character in 'The Call of the Wild'?",
+ # "answer": "Buck",
+ # },
+ # {"question": "Who wrote 'The Call of the Wild'?", "answer": "Jack London"},
+ # {
+ # "question": "Where does Buck live at the start of the book?",
+ # "answer": "In the Santa Clara Valley, at Judge Miller’s place.",
+ # },
+ # {
+ # "question": "Why is Buck kidnapped?",
+ # "answer": "He is kidnapped to be sold as a sled dog in the Yukon during the Klondike Gold Rush.",
+ # },
+ # {
+ # "question": "How does Buck become the leader of the sled dog team?",
+ # "answer": "Buck becomes the leader after defeating the original leader, Spitz, in a fight.",
+ # },
+ # ]
# "https://www.ibiblio.org/ebooks/London/Call%20of%20Wild.pdf"
- # http://public-library.uk/ebooks/59/83.pdf
+ # # http://public-library.uk/ebooks/59/83.pdf
# result = await start_test(
- # ".data/3ZCCCW.pdf",
+ # [".data/3ZCCCW.pdf"],
# test_set=test_set,
# user_id="677",
# params=["chunk_size", "search_type"],
@@ -739,7 +755,7 @@ async def main():
# )
parser = argparse.ArgumentParser(description="Run tests against a document.")
- parser.add_argument("--file", required=True, help="URL or location of the document to test.")
+ parser.add_argument("--file", nargs="+", required=True, help="List of file paths to test.")
parser.add_argument("--test_set", required=True, help="Path to JSON file containing the test set.")
parser.add_argument("--user_id", required=True, help="User ID.")
parser.add_argument("--params", help="Additional parameters in JSON format.")
@@ -776,6 +792,7 @@ async def main():
return
else:
params = None
+ logging.info("Args datatype is", type(args.file))
#clean up params here
await start_test(data=args.file, test_set=test_set, user_id= args.user_id, params= params, metadata =metadata, retriever_type=args.retriever_type)
diff --git a/level_3/scripts/create_database.py b/level_3/scripts/create_database.py
index b06021ddd..f6715741b 100644
--- a/level_3/scripts/create_database.py
+++ b/level_3/scripts/create_database.py
@@ -15,6 +15,7 @@ import models.sessions
import models.testoutput
import models.testset
import models.user
+import models.docs
from sqlalchemy import create_engine, text
import psycopg2
from dotenv import load_dotenv
diff --git a/level_3/vectordb/basevectordb.py b/level_3/vectordb/basevectordb.py
index 69ce66784..81f4f7618 100644
--- a/level_3/vectordb/basevectordb.py
+++ b/level_3/vectordb/basevectordb.py
@@ -289,3 +289,9 @@ class BaseMemory:
async def delete_memories(self, namespace:str, params: Optional[str] = None):
return await self.vector_db.delete_memories(namespace,params)
+
+ async def count_memories(self, namespace:str, params: Optional[str] = None):
+ return await self.vector_db.count_memories(namespace,params)
+
+
+
diff --git a/level_3/vectordb/loaders/loaders.py b/level_3/vectordb/loaders/loaders.py
index 2ed80fe2b..acb54147e 100644
--- a/level_3/vectordb/loaders/loaders.py
+++ b/level_3/vectordb/loaders/loaders.py
@@ -6,50 +6,72 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from vectordb.chunkers.chunkers import chunk_data
from llama_hub.file.base import SimpleDirectoryReader
-
+from langchain.document_loaders import UnstructuredURLLoader
from langchain.document_loaders import DirectoryLoader
-
+import logging
+import os
+from langchain.document_loaders import TextLoader
import requests
async def _document_loader( observation: str, loader_settings: dict):
- # Check the format of the document
+
document_format = loader_settings.get("format", "text")
loader_strategy = loader_settings.get("strategy", "VANILLA")
chunk_size = loader_settings.get("chunk_size", 500)
chunk_overlap = loader_settings.get("chunk_overlap", 20)
+ import logging
+ import os
- print("LOADER SETTINGS", loader_settings)
+ logging.info("LOADER SETTINGS %s", loader_settings)
- if document_format == "PDF":
- if loader_settings.get("source") == "URL":
- pdf_response = requests.get(loader_settings["path"])
- pdf_stream = BytesIO(pdf_response.content)
- with fitz.open(stream=pdf_stream, filetype='pdf') as doc:
- file_content = ""
- for page in doc:
- file_content += page.get_text()
- pages = chunk_data(chunk_strategy= loader_strategy, source_data=file_content, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+ list_of_docs = loader_settings["path"]
+ chunked_doc = []
- return pages
- elif loader_settings.get("source") == "DEVICE":
- import os
+ if loader_settings.get("source") == "URL":
+ for file in list_of_docs:
+ if document_format == "PDF":
+ pdf_response = requests.get(file)
+ pdf_stream = BytesIO(pdf_response.content)
+ with fitz.open(stream=pdf_stream, filetype='pdf') as doc:
+ file_content = ""
+ for page in doc:
+ file_content += page.get_text()
+ pages = chunk_data(chunk_strategy=loader_strategy, source_data=file_content, chunk_size=chunk_size,
+ chunk_overlap=chunk_overlap)
- current_directory = os.getcwd()
- import logging
- logging.info("Current Directory: %s", current_directory)
+ chunked_doc.append(pages)
- loader = DirectoryLoader(".data", recursive=True)
+ elif document_format == "TEXT":
+ loader = UnstructuredURLLoader(urls=file)
+ file_content = loader.load()
+ pages = chunk_data(chunk_strategy=loader_strategy, source_data=file_content, chunk_size=chunk_size,
+ chunk_overlap=chunk_overlap)
+ chunked_doc.append(pages)
+ elif loader_settings.get("source") == "DEVICE":
+
+ current_directory = os.getcwd()
+ logging.info("Current Directory: %s", current_directory)
+
+ loader = DirectoryLoader(".data", recursive=True)
+ if document_format == "PDF":
# loader = SimpleDirectoryReader(".data", recursive=True, exclude_hidden=True)
documents = loader.load()
logging.info("Documents: %s", documents)
# pages = documents.load_and_split()
- return documents
+ chunked_doc.append(documents)
- elif document_format == "TEXT":
- pages = chunk_data(chunk_strategy= loader_strategy, source_data=observation, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
- return pages
+
+ elif document_format == "TEXT":
+ documents = loader.load()
+ logging.info("Documents: %s", documents)
+ # pages = documents.load_and_split()
+ chunked_doc.append(documents)
else:
- raise ValueError(f"Unsupported document format: {document_format}")
+ raise ValueError(f"Error: ")
+ return chunked_doc
+
+
+
diff --git a/level_3/vectordb/vectordb.py b/level_3/vectordb/vectordb.py
index 55894a405..0cc66066b 100644
--- a/level_3/vectordb/vectordb.py
+++ b/level_3/vectordb/vectordb.py
@@ -153,7 +153,7 @@ class WeaviateVectorDB(VectorDB):
# Assuming _document_loader returns a list of documents
documents = await _document_loader(observation, loader_settings)
logging.info("here are the docs %s", str(documents))
- for doc in documents:
+ for doc in documents[0]:
document_to_load = self._stuct(doc.page_content, params, metadata_schema_class)
logging.info("Loading document with provided loader settings %s", str(document_to_load))
@@ -290,6 +290,30 @@ class WeaviateVectorDB(VectorDB):
},
)
+
+ async def count_memories(self, namespace: str = None, params: dict = None) -> int:
+ """
+ Count memories in a Weaviate database.
+
+ Args:
+ namespace (str, optional): The Weaviate namespace to count memories in. If not provided, uses the default namespace.
+
+ Returns:
+ int: The number of memories in the specified namespace.
+ """
+ if namespace is None:
+ namespace = self.namespace
+
+ client = self.init_weaviate(namespace =namespace)
+
+ try:
+ object_count = client.query.aggregate(namespace).with_meta_count().do()
+ return object_count
+ except Exception as e:
+ logging.info(f"Error counting memories: {str(e)}")
+ # Handle the error or log it
+ return 0
+
def update_memories(self, observation, namespace: str, params: dict = None):
client = self.init_weaviate(namespace = self.namespace)
diff --git a/level_3/wait-for-it.sh b/level_3/wait-for-it.sh
new file mode 100644
index 000000000..3974640b0
--- /dev/null
+++ b/level_3/wait-for-it.sh
@@ -0,0 +1,182 @@
+#!/usr/bin/env bash
+# Use this script to test if a given TCP host/port are available
+
+WAITFORIT_cmdname=${0##*/}
+
+echoerr() { if [[ $WAITFORIT_QUIET -ne 1 ]]; then echo "$@" 1>&2; fi }
+
+usage()
+{
+ cat << USAGE >&2
+Usage:
+ $WAITFORIT_cmdname host:port [-s] [-t timeout] [-- command args]
+ -h HOST | --host=HOST Host or IP under test
+ -p PORT | --port=PORT TCP port under test
+ Alternatively, you specify the host and port as host:port
+ -s | --strict Only execute subcommand if the test succeeds
+ -q | --quiet Don't output any status messages
+ -t TIMEOUT | --timeout=TIMEOUT
+ Timeout in seconds, zero for no timeout
+ -- COMMAND ARGS Execute command with args after the test finishes
+USAGE
+ exit 1
+}
+
+wait_for()
+{
+ if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then
+ echoerr "$WAITFORIT_cmdname: waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT"
+ else
+ echoerr "$WAITFORIT_cmdname: waiting for $WAITFORIT_HOST:$WAITFORIT_PORT without a timeout"
+ fi
+ WAITFORIT_start_ts=$(date +%s)
+ while :
+ do
+ if [[ $WAITFORIT_ISBUSY -eq 1 ]]; then
+ nc -z $WAITFORIT_HOST $WAITFORIT_PORT
+ WAITFORIT_result=$?
+ else
+ (echo -n > /dev/tcp/$WAITFORIT_HOST/$WAITFORIT_PORT) >/dev/null 2>&1
+ WAITFORIT_result=$?
+ fi
+ if [[ $WAITFORIT_result -eq 0 ]]; then
+ WAITFORIT_end_ts=$(date +%s)
+ echoerr "$WAITFORIT_cmdname: $WAITFORIT_HOST:$WAITFORIT_PORT is available after $((WAITFORIT_end_ts - WAITFORIT_start_ts)) seconds"
+ break
+ fi
+ sleep 1
+ done
+ return $WAITFORIT_result
+}
+
+wait_for_wrapper()
+{
+ # In order to support SIGINT during timeout: http://unix.stackexchange.com/a/57692
+ if [[ $WAITFORIT_QUIET -eq 1 ]]; then
+ timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --quiet --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT &
+ else
+ timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT &
+ fi
+ WAITFORIT_PID=$!
+ trap "kill -INT -$WAITFORIT_PID" INT
+ wait $WAITFORIT_PID
+ WAITFORIT_RESULT=$?
+ if [[ $WAITFORIT_RESULT -ne 0 ]]; then
+ echoerr "$WAITFORIT_cmdname: timeout occurred after waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT"
+ fi
+ return $WAITFORIT_RESULT
+}
+
+# process arguments
+while [[ $# -gt 0 ]]
+do
+ case "$1" in
+ *:* )
+ WAITFORIT_hostport=(${1//:/ })
+ WAITFORIT_HOST=${WAITFORIT_hostport[0]}
+ WAITFORIT_PORT=${WAITFORIT_hostport[1]}
+ shift 1
+ ;;
+ --child)
+ WAITFORIT_CHILD=1
+ shift 1
+ ;;
+ -q | --quiet)
+ WAITFORIT_QUIET=1
+ shift 1
+ ;;
+ -s | --strict)
+ WAITFORIT_STRICT=1
+ shift 1
+ ;;
+ -h)
+ WAITFORIT_HOST="$2"
+ if [[ $WAITFORIT_HOST == "" ]]; then break; fi
+ shift 2
+ ;;
+ --host=*)
+ WAITFORIT_HOST="${1#*=}"
+ shift 1
+ ;;
+ -p)
+ WAITFORIT_PORT="$2"
+ if [[ $WAITFORIT_PORT == "" ]]; then break; fi
+ shift 2
+ ;;
+ --port=*)
+ WAITFORIT_PORT="${1#*=}"
+ shift 1
+ ;;
+ -t)
+ WAITFORIT_TIMEOUT="$2"
+ if [[ $WAITFORIT_TIMEOUT == "" ]]; then break; fi
+ shift 2
+ ;;
+ --timeout=*)
+ WAITFORIT_TIMEOUT="${1#*=}"
+ shift 1
+ ;;
+ --)
+ shift
+ WAITFORIT_CLI=("$@")
+ break
+ ;;
+ --help)
+ usage
+ ;;
+ *)
+ echoerr "Unknown argument: $1"
+ usage
+ ;;
+ esac
+done
+
+if [[ "$WAITFORIT_HOST" == "" || "$WAITFORIT_PORT" == "" ]]; then
+ echoerr "Error: you need to provide a host and port to test."
+ usage
+fi
+
+WAITFORIT_TIMEOUT=${WAITFORIT_TIMEOUT:-15}
+WAITFORIT_STRICT=${WAITFORIT_STRICT:-0}
+WAITFORIT_CHILD=${WAITFORIT_CHILD:-0}
+WAITFORIT_QUIET=${WAITFORIT_QUIET:-0}
+
+# Check to see if timeout is from busybox?
+WAITFORIT_TIMEOUT_PATH=$(type -p timeout)
+WAITFORIT_TIMEOUT_PATH=$(realpath $WAITFORIT_TIMEOUT_PATH 2>/dev/null || readlink -f $WAITFORIT_TIMEOUT_PATH)
+
+WAITFORIT_BUSYTIMEFLAG=""
+if [[ $WAITFORIT_TIMEOUT_PATH =~ "busybox" ]]; then
+ WAITFORIT_ISBUSY=1
+ # Check if busybox timeout uses -t flag
+ # (recent Alpine versions don't support -t anymore)
+ if timeout &>/dev/stdout | grep -q -e '-t '; then
+ WAITFORIT_BUSYTIMEFLAG="-t"
+ fi
+else
+ WAITFORIT_ISBUSY=0
+fi
+
+if [[ $WAITFORIT_CHILD -gt 0 ]]; then
+ wait_for
+ WAITFORIT_RESULT=$?
+ exit $WAITFORIT_RESULT
+else
+ if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then
+ wait_for_wrapper
+ WAITFORIT_RESULT=$?
+ else
+ wait_for
+ WAITFORIT_RESULT=$?
+ fi
+fi
+
+if [[ $WAITFORIT_CLI != "" ]]; then
+ if [[ $WAITFORIT_RESULT -ne 0 && $WAITFORIT_STRICT -eq 1 ]]; then
+ echoerr "$WAITFORIT_cmdname: strict mode, refusing to execute subprocess"
+ exit $WAITFORIT_RESULT
+ fi
+ exec "${WAITFORIT_CLI[@]}"
+else
+ exit $WAITFORIT_RESULT
+fi
\ No newline at end of file