diff --git a/core/local_model.py b/core/local_model.py index e00b52d..5fc5b16 100644 --- a/core/local_model.py +++ b/core/local_model.py @@ -150,7 +150,7 @@ def simple_qa(query: str, context: str) -> str: temperature=0.1, messages=messages, stream=False, - response_format="web", + response_format={"type": "text"}, # mlc doesn't supports string "web" ) return response.choices[0].message.content diff --git a/core/tasks/celery_app.py b/core/tasks/celery_app.py new file mode 100644 index 0000000..4e6922f --- /dev/null +++ b/core/tasks/celery_app.py @@ -0,0 +1,10 @@ +import redis +from celery import Celery +from typing import cast +import core.config as configs + + +redis_client = cast(redis.Redis, redis.Redis.from_url(configs.redis_url)) # annoyingly from_url returns None, not Self +app = Celery("tasks", broker=configs.redis_url) + +app.autodiscover_tasks(["core.tasks"]) # Explicitly discover tasks in 'app' package diff --git a/core/tasks/is_ready.py b/core/tasks/is_ready.py index fdf9a04..ae59bcf 100644 --- a/core/tasks/is_ready.py +++ b/core/tasks/is_ready.py @@ -1,11 +1,11 @@ -from .tasks import app - import socket import requests from core.config import litellm_url, vector_db_url +from .celery_app import app + def is_ready(): # response = requests.get(f"{litellm_url}/health", headers={ # "Authorization": f"Bearer {os.getenv('LITELLM_MASTER_KEY', '')}" diff --git a/core/tasks/tasks.py b/core/tasks/tasks.py index 8b4129f..a58578d 100644 --- a/core/tasks/tasks.py +++ b/core/tasks/tasks.py @@ -44,18 +44,15 @@ import core.config as configs from .is_ready import is_ready - -redis_client = cast(redis.Redis, redis.Redis.from_url(configs.redis_url)) # annoyingly from_url returns None, not Self -app = Celery("tasks", broker=configs.redis_url) - -app.autodiscover_tasks(["core.tasks"]) # Explicitly discover tasks in 'app' package +from .celery_app import app # Global MongoDB client -mongo_client: MongoClient +mongo_client: MongoClient = None +redis_client = cast(redis.Redis, redis.Redis.from_url(configs.redis_url)) # annoyingly from_url returns None, not Self @signals.worker_process_init.connect -async def ensure_connections(*args, **kwargs): +def ensure_connections(*args, **kwargs): global mongo_client mongo_client = MongoClient(configs.mongo_url) @@ -220,11 +217,12 @@ def form_openai_tools(tools, assistant_id: str): @shared_task def execute_chat_completion(assistant_id, thread_id, redis_channel, run_id): try: + db = mongo_client[configs.mongo_database] # OpenAI call can fail, so we need to get the db again + oai_client = OpenAI( base_url=configs.litellm_url, api_key=os.getenv("LITELLM_MASTER_KEY"), # point to litellm server ) - db = mongo_client[configs.mongo_database] # Fetch assistant and thread messages synchronously assistant = db.assistants.find_one({"id": assistant_id}) @@ -459,6 +457,8 @@ def execute_chat_completion(assistant_id, thread_id, redis_channel, run_id): @app.task def execute_asst_file_create(file_id: str, assistant_id: str): try: + if mongo_client is None: + raise Exception("MongoDB client not initialized yet") db = mongo_client[configs.mongo_database] collection_name = assistant_id text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0) diff --git a/core/tools/knowledge/file_knowledge_tool.py b/core/tools/knowledge/file_knowledge_tool.py index 105cb11..dea2088 100644 --- a/core/tools/knowledge/file_knowledge_tool.py +++ b/core/tools/knowledge/file_knowledge_tool.py @@ -6,7 +6,7 @@ # Third Party import requests -vector_db_url = f"{configs.vector_db_url}/similarity_search" +vector_db_url = f"{configs.vector_db_url}/similarity_match" class FileKnowledgeTool: name = "FileKnowledge" diff --git a/core/tools/knowledge/vector_db/milvus/query_milvus.py b/core/tools/knowledge/vector_db/milvus/query_milvus.py index 35b6677..725cbca 100644 --- a/core/tools/knowledge/vector_db/milvus/query_milvus.py +++ b/core/tools/knowledge/vector_db/milvus/query_milvus.py @@ -133,6 +133,7 @@ def __init__( if alias is not None: self.alias = alias elif connection_args is not None: + connection_args = DEFAULT_MILVUS_CONNECTION self.alias = Milvus.create_connection_alias(connection_args) else: raise ValueError('alias or connection_args must be passed to Milvus construtor') diff --git a/docker-compose.yml b/docker-compose.yml index ea091a7..01c74d2 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -115,6 +115,7 @@ services: - EMBEDDING_HOST=text-embedding-api - VECTOR_DB_HOST=vector-db-api - MILVUS_HOST=milvus + - LITELLM_MASTER_KEY=abc depends_on: - redis - mongodb @@ -134,6 +135,7 @@ services: - REDIS_HOST=redis - MONGODB_HOST=mongodb - LITELLM_HOST=litellm + - MILVUS_HOST=milvus ports: - '8000:8000' depends_on: diff --git a/services/backend/api_server/app/backend.py b/services/backend/api_server/app/backend.py index 11ed38d..3fdab5c 100644 --- a/services/backend/api_server/app/backend.py +++ b/services/backend/api_server/app/backend.py @@ -127,16 +127,16 @@ # MongoDB Configurationget LITELLM_URL = configs.litellm_url -LITELLM_MASTER_KEY = os.getenv("LITELLM_MASTER_KEY", "") +LITELLM_MASTER_KEY = os.getenv("LITELLM_MASTER_KEY", "abc") # Litellm fails without this key HEADERS = {"accept": "application/json", "Content-Type": "application/json"} # Initialize MongoDB client mongo_client = AsyncIOMotorClient(configs.mongo_url, server_api=ServerApi("1")) database = mongo_client[configs.mongo_database] -celery_app = Celery(configs.redis_url) +celery_app = Celery(broker=configs.redis_url) -redis = aioredis.from_url(configs.redis_url) +redis = aioredis.from_url(configs.redis_url, encoding="utf-8", decode_responses=True) logging.basicConfig(level=logging.INFO) @@ -764,6 +764,14 @@ async def redis_subscriber(channel, timeout=1): pubsub = redis.pubsub() await pubsub.subscribe(channel) + # Check if the subscription was successful + channels = await redis.pubsub_channels() + logging.info(f"Channels: {channels}") + if channel in channels: + logging.info(f"Successfully subscribed to channel: {channel}") + else: + logging.error(f"Failed to subscribe to channel: {channel}") + while True: try: message = await asyncio.wait_for(