update

app/llmops/components/get_documents/run.py

@@ -16,7 +16,7 @@ def go(args):
     zip_path = os.path.join(args.path_document_folder, f"{args.document_folder}.zip")
     shutil.make_archive(zip_path.replace('.zip', ''), 'zip', args.path_document_folder, args.document_folder)
 
-    with mlflow.start_run(experiment_id=mlflow.get_experiment_by_name("development").experiment_id) as run:
+    with mlflow.start_run(experiment_id=mlflow.get_experiment_by_name("development").experiment_id):
 
         existing_params = mlflow.get_run(mlflow.active_run().info.run_id).data.params
         if 'artifact_description' not in existing_params:

app/llmops/components/wandb_utils/log_artifact.py

@@ -1,6 +1,4 @@
 import wandb
-import mlflow
-
 
 def log_artifact(artifact_name, artifact_type, artifact_description, filename, wandb_run):
     """

app/llmops/main.py

@@ -1,11 +1,8 @@
-import json
-
 import mlflow
 import tempfile
 import os
 import hydra
 from omegaconf import DictConfig
-from decouple import config
 
 _steps = [
     "get_documents",
@@ -27,7 +24,7 @@ def go(config: DictConfig):
     active_steps = steps_par.split(",") if steps_par != "all" else _steps
 
     # Move to a temporary directory
-    with tempfile.TemporaryDirectory() as tmp_dir:
+    with tempfile.TemporaryDirectory():
 
         if "get_documents" in active_steps:
             # Download file and load in W&B

app/llmops/src/etl_chromadb_scanned_pdf/run.py

@@ -10,8 +10,6 @@ import shutil
 
 import chromadb
 # from openai import OpenAI
-from typing import List
-import numpy as np
 import pytesseract as pt
 from pdf2image import convert_from_path
 from langchain.schema import Document

app/llmops/tests/src/chain_of_thought/test_run.py

@@ -1,31 +0,0 @@
-import pytest
-from unittest.mock import patch, MagicMock
-import sys
-sys.path.append("/Users/leehongkai/projects/aimingmed/aimingmed-ai/app/llmops")
-from src.chain_of_thought import run
-
-def test_go():
-    # Create mock arguments
-    args = MagicMock()
-    args.query = "test_query"
-    args.input_chromadb_artifact = "test_artifact"
-    args.embedding_model = "test_embedding_model"
-    args.chat_model_provider = "gemini"
-
-    # Mock wandb.init and other external dependencies
-    with patch("wandb.init") as mock_wandb_init, \
-         patch("chromadb.PersistentClient") as mock_chromadb_client, \
-         patch("sentence_transformers.SentenceTransformer") as mock_sentence_transformer, \
-         patch("langchain_google_genai.ChatGoogleGenerativeAI") as mock_chat_google_generative_ai:
-
-        # Configure the mocks
-        mock_wandb_init.return_value = MagicMock()
-        mock_chromadb_client.return_value = MagicMock()
-        mock_sentence_transformer.return_value = MagicMock()
-        mock_chat_google_generative_ai.return_value = MagicMock()
-
-        # Call the go function
-        run.go(args)
-
-        # Add assertions to validate the behavior of the go function
-        assert mock_wandb_init.called

app/streamlit/tests/test_chatbot.py

@@ -1,7 +1,5 @@
-import os
 import pytest
 import chromadb
-from langchain.prompts import PromptTemplate
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_deepseek import ChatDeepSeek
 from langchain_community.llms.moonshot import Moonshot