From cb8927d87bcfa5513b46c3e6d089c9d2ed764362 Mon Sep 17 00:00:00 2001 From: User Date: Mon, 27 May 2024 15:07:49 +0300 Subject: [PATCH] replace logging calls with logger calls --- examples/_test_single_llama_index.py | 2 +- examples/_test_single_openai_tools_react.py | 2 +- examples/delegation_crewai.py | 2 +- examples/delegation_demo.py | 2 +- examples/image_generation_crewai.py | 2 +- examples/math_crewai.py | 2 +- .../research_agent/research_agent_main.py | 2 +- examples/test_single_crewai_agent.py | 2 +- motleycrew/agents/parent.py | 12 ++++---- .../research_agent/answer_task.py | 4 +-- .../research_agent/question_generator.py | 6 ++-- .../research_agent/question_task.py | 8 ++--- motleycrew/crew.py | 24 +++++++-------- motleycrew/storage/graph_store_utils.py | 7 ++--- motleycrew/storage/kuzu_graph_store.py | 30 +++++++++---------- motleycrew/tasks/simple.py | 8 ++--- motleycrew/tools/image/dall_e.py | 14 ++++----- motleycrew/tracking/callbacks.py | 6 ++-- motleycrew/tracking/utils.py | 6 ++-- tests/run_integration_tests.py | 19 ++++++------ 20 files changed, 79 insertions(+), 81 deletions(-) diff --git a/examples/_test_single_llama_index.py b/examples/_test_single_llama_index.py index 454916ea..1c0db26d 100644 --- a/examples/_test_single_llama_index.py +++ b/examples/_test_single_llama_index.py @@ -4,7 +4,7 @@ from motleycrew import MotleyCrew from motleycrew.agents.llama_index import ReActLlamaIndexMotleyAgent -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging from motleycrew.tasks import SimpleTask diff --git a/examples/_test_single_openai_tools_react.py b/examples/_test_single_openai_tools_react.py index 01837c96..af6f2518 100644 --- a/examples/_test_single_openai_tools_react.py +++ b/examples/_test_single_openai_tools_react.py @@ -4,7 +4,7 @@ from motleycrew import MotleyCrew from motleycrew.agents.langchain.openai_tools_react import ReactOpenAIToolsAgent from motleycrew.agents.langchain.react import ReactMotleyAgent -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging from motleycrew.tasks import SimpleTask from motleycache import enable_cache diff --git a/examples/delegation_crewai.py b/examples/delegation_crewai.py index aa5f7358..5a523023 100644 --- a/examples/delegation_crewai.py +++ b/examples/delegation_crewai.py @@ -12,7 +12,7 @@ from motleycrew import MotleyCrew from motleycrew.agents.crewai import CrewAIMotleyAgent from motleycrew.agents.langchain.react import ReactMotleyAgent -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging from motleycrew.tasks import SimpleTask WORKING_DIR = Path(os.path.realpath(".")) diff --git a/examples/delegation_demo.py b/examples/delegation_demo.py index 09f69fc3..ba4a4fae 100644 --- a/examples/delegation_demo.py +++ b/examples/delegation_demo.py @@ -13,7 +13,7 @@ from motleycrew.agents.langchain.react import ReactMotleyAgent from motleycrew.agents.llama_index import ReActLlamaIndexMotleyAgent from motleycrew.tools.image.dall_e import DallEImageGeneratorTool -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging from motleycrew.tasks import SimpleTask WORKING_DIR = Path(os.path.realpath(".")) diff --git a/examples/image_generation_crewai.py b/examples/image_generation_crewai.py index b79a76ea..c6721d59 100644 --- a/examples/image_generation_crewai.py +++ b/examples/image_generation_crewai.py @@ -3,7 +3,7 @@ from motleycrew import MotleyCrew from motleycrew.agents.crewai import CrewAIMotleyAgent from motleycrew.tools.image.dall_e import DallEImageGeneratorTool -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging def main(): diff --git a/examples/math_crewai.py b/examples/math_crewai.py index 58c31692..51b093b8 100644 --- a/examples/math_crewai.py +++ b/examples/math_crewai.py @@ -3,7 +3,7 @@ from motleycrew import MotleyCrew from motleycrew.agents.crewai import CrewAIMotleyAgent from motleycrew.tools import PythonREPLTool -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging def main(): diff --git a/examples/research_agent/research_agent_main.py b/examples/research_agent/research_agent_main.py index 72df7891..e003701c 100644 --- a/examples/research_agent/research_agent_main.py +++ b/examples/research_agent/research_agent_main.py @@ -8,7 +8,7 @@ from motleycrew import MotleyCrew from motleycrew.storage import MotleyKuzuGraphStore -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging from motleycrew.applications.research_agent.question_task import QuestionTask from motleycrew.applications.research_agent.answer_task import AnswerTask diff --git a/examples/test_single_crewai_agent.py b/examples/test_single_crewai_agent.py index 03bce351..c1356e3e 100644 --- a/examples/test_single_crewai_agent.py +++ b/examples/test_single_crewai_agent.py @@ -3,7 +3,7 @@ from motleycrew import MotleyCrew from motleycrew.agents.crewai import CrewAIMotleyAgent -from motleycrew.common.utils import configure_logging +from motleycrew.common import configure_logging def main(): diff --git a/motleycrew/agents/parent.py b/motleycrew/agents/parent.py index eb7df005..65d48785 100644 --- a/motleycrew/agents/parent.py +++ b/motleycrew/agents/parent.py @@ -1,4 +1,3 @@ -import logging from typing import TYPE_CHECKING, Optional, Sequence from langchain_core.tools import Tool @@ -9,6 +8,7 @@ from motleycrew.tools import MotleyTool from motleycrew.common import MotleyAgentFactory, MotleySupportedTool from motleycrew.common.exceptions import AgentNotMaterialized, CannotModifyMaterializedAgent +from motleycrew.common import logger if TYPE_CHECKING: from motleycrew import MotleyCrew @@ -57,7 +57,7 @@ def is_materialized(self): def materialize(self): if self.is_materialized: - logging.info("Agent is already materialized, skipping materialization") + logger.info("Agent is already materialized, skipping materialization") return assert self.agent_factory, "Cannot materialize agent without a factory provided" self._agent = self.agent_factory(tools=self.tools) @@ -91,7 +91,7 @@ def call_agent(*args, **kwargs): ) # def call_as_tool(self, *args, **kwargs) -> Any: - # logging.info("Entering delegation for %s", self.name) + # logger.info("Entering delegation for %s", self.name) # assert self.crew, "can't accept delegated task outside of a crew" # # if len(args) > 0: @@ -102,7 +102,7 @@ def call_agent(*args, **kwargs): # else: # input_ = json.dumps(kwargs) # - # logging.info("Made the args: %s", input_) + # logger.info("Made the args: %s", input_) # # # TODO: pass context of parent task to agent nicely? # # TODO: mark the current task as depending on the new task @@ -117,11 +117,11 @@ def call_agent(*args, **kwargs): # ) # # # TODO: make sure tools return task objects, which are properly used by callers - # logging.info("Executing subtask '%s'", task.name) + # logger.info("Executing subtask '%s'", task.name) # self.crew.task_graph.set_task_running(task=task) # result = self.crew.execute(task, return_result=True) # - # logging.info("Finished subtask '%s' - %s", task.name, result) + # logger.info("Finished subtask '%s' - %s", task.name, result) # self.crew.task_graph.set_task_done(task=task) # # return result diff --git a/motleycrew/applications/research_agent/answer_task.py b/motleycrew/applications/research_agent/answer_task.py index 96cff2e4..f3857192 100644 --- a/motleycrew/applications/research_agent/answer_task.py +++ b/motleycrew/applications/research_agent/answer_task.py @@ -1,4 +1,3 @@ -import logging from typing import List, Optional from langchain_core.runnables import Runnable @@ -10,6 +9,7 @@ from motleycrew.applications.research_agent.question import Question, QuestionAnsweringTaskUnit from motleycrew.applications.research_agent.question_answerer import AnswerSubQuestionTool from motleycrew.storage import MotleyGraphStore +from motleycrew.common import logger class AnswerTask(Task): @@ -34,7 +34,7 @@ def get_next_unit(self) -> QuestionAnsweringTaskUnit | None: ).format(Question.get_label(), Question.get_label()) query_result = self.graph_store.run_cypher_query(query, container=Question) - logging.info("Available questions: %s", query_result) + logger.info("Available questions: %s", query_result) if not query_result: return None else: diff --git a/motleycrew/applications/research_agent/question_generator.py b/motleycrew/applications/research_agent/question_generator.py index 1dae75e3..96466626 100644 --- a/motleycrew/applications/research_agent/question_generator.py +++ b/motleycrew/applications/research_agent/question_generator.py @@ -1,6 +1,5 @@ from typing import Optional from pathlib import Path -import logging from langchain_core.language_models import BaseLanguageModel from langchain_core.runnables import ( @@ -18,6 +17,7 @@ from motleycrew.common.llms import init_llm from motleycrew.common.utils import print_passthrough from motleycrew.storage import MotleyGraphStore +from motleycrew.common import logger from motleycrew.applications.research_agent.question import Question, QuestionGenerationTaskUnit @@ -103,10 +103,10 @@ def insert_questions(input_dict) -> None: questions_raw = input_dict["subquestions"].content questions = [q.strip() for q in questions_raw.split("\n") if len(q.strip()) > 1] for q in questions: - logging.info("Inserting question: %s", q) + logger.info("Inserting question: %s", q) subquestion = graph.insert_node(Question(question=q)) graph.create_relation(input_dict["question"], subquestion, IS_SUBQUESTION_PREDICATE) - logging.info("Inserted %s questions", len(questions)) + logger.info("Inserted %s questions", len(questions)) def set_context(input_dict: dict): node = input_dict["question"] diff --git a/motleycrew/applications/research_agent/question_task.py b/motleycrew/applications/research_agent/question_task.py index 044303c9..ca7091de 100644 --- a/motleycrew/applications/research_agent/question_task.py +++ b/motleycrew/applications/research_agent/question_task.py @@ -1,4 +1,3 @@ -import logging from typing import List, Optional from langchain_core.runnables import Runnable @@ -10,6 +9,7 @@ from .question import Question, QuestionGenerationTaskUnit from .question_generator import QuestionGeneratorTool from .question_prioritizer import QuestionPrioritizerTool +from motleycrew.common import logger class QuestionTask(Task): @@ -39,7 +39,7 @@ def get_next_unit(self) -> QuestionGenerationTaskUnit | None: return None unanswered_questions = self.get_unanswered_questions(only_without_children=True) - logging.info("Loaded unanswered questions: %s", unanswered_questions) + logger.info("Loaded unanswered questions: %s", unanswered_questions) if not len(unanswered_questions): return None @@ -50,11 +50,11 @@ def get_next_unit(self) -> QuestionGenerationTaskUnit | None: "unanswered_questions": unanswered_questions, } ) - logging.info("Most pertinent question according to the tool: %s", most_pertinent_question) + logger.info("Most pertinent question according to the tool: %s", most_pertinent_question) return QuestionGenerationTaskUnit(question=most_pertinent_question) def register_completed_unit(self, unit: TaskUnitType) -> None: - logging.info("==== Completed iteration %s of %s ====", self.n_iter + 1, self.max_iter) + logger.info("==== Completed iteration %s of %s ====", self.n_iter + 1, self.max_iter) self.n_iter += 1 if self.n_iter >= self.max_iter: self.set_done(True) diff --git a/motleycrew/crew.py b/motleycrew/crew.py index a0b0f2af..291027dd 100644 --- a/motleycrew/crew.py +++ b/motleycrew/crew.py @@ -1,5 +1,4 @@ from typing import Collection, Sequence, Optional -import logging import os from motleycrew.agents.parent import MotleyAgentParent @@ -7,6 +6,7 @@ from motleycrew.storage import MotleyGraphStore from motleycrew.storage.graph_store_utils import init_graph_store from motleycrew.tools import MotleyTool +from motleycrew.common import logger class MotleyCrew: @@ -39,7 +39,7 @@ def create_simple_task( def run(self) -> list[TaskUnit]: if not self.single_thread: - logging.warning("Multithreading is not implemented yet, will run in single thread") + logger.warning("Multithreading is not implemented yet, will run in single thread") return self._run_sync() @@ -70,24 +70,24 @@ def _run_sync(self) -> list[TaskUnit]: did_something = False available_tasks = self.get_available_tasks() - logging.info("Available tasks: %s", available_tasks) + logger.info("Available tasks: %s", available_tasks) for task in available_tasks: - logging.info("Processing task: %s", task) + logger.info("Processing task: %s", task) next_unit = task.get_next_unit() if next_unit is None: - logging.info("Got no matching units for task %s", task) + logger.info("Got no matching units for task %s", task) else: - logging.info("Got a matching unit for task %s", task) + logger.info("Got a matching unit for task %s", task) current_unit = next_unit - logging.info("Processing task: %s", current_unit) + logger.info("Processing task: %s", current_unit) extra_tools = self.get_extra_tools(task) agent = task.get_worker(extra_tools) - logging.info("Assigned unit %s to agent %s, dispatching", current_unit, agent) + logger.info("Assigned unit %s to agent %s, dispatching", current_unit, agent) current_unit.set_running() task.register_started_unit(current_unit) @@ -95,7 +95,7 @@ def _run_sync(self) -> list[TaskUnit]: result = agent.invoke(current_unit.as_dict()) current_unit.output = result - logging.info("Task unit %s completed, marking as done", current_unit) + logger.info("Task unit %s completed, marking as done", current_unit) current_unit.set_done() task.register_completed_unit(current_unit) done_units.append(current_unit) @@ -104,7 +104,7 @@ def _run_sync(self) -> list[TaskUnit]: continue if not did_something: - logging.info("Nothing left to do, exiting") + logger.info("Nothing left to do, exiting") return done_units def get_available_tasks(self) -> list[Task]: @@ -138,7 +138,7 @@ def get_available_tasks(self) -> list[Task]: # raise exc # # task = future.mc_task - # logging.info(f"Finished task '{task.name}'") + # logger.info(f"Finished task '{task.name}'") # self.futures.remove(future) # tasks.set_task_done(task) # self.adispatch_next_batch() @@ -150,7 +150,7 @@ def get_available_tasks(self) -> list[Task]: # next_ = self.task_graph.get_ready_tasks() # for t in next_: # self.task_graph.set_task_running(t) - # logging.info(f"Dispatching task '{t.name}'") + # logger.info(f"Dispatching task '{t.name}'") # future = self.thread_pool.submit( # self.execute, # t, diff --git a/motleycrew/storage/graph_store_utils.py b/motleycrew/storage/graph_store_utils.py index a64e9fb2..acce5927 100644 --- a/motleycrew/storage/graph_store_utils.py +++ b/motleycrew/storage/graph_store_utils.py @@ -1,11 +1,10 @@ import tempfile from typing import Optional -import logging import os from motleycrew.common import Defaults from motleycrew.common import GraphStoreType - +from motleycrew.common import logger from motleycrew.storage import MotleyKuzuGraphStore @@ -17,10 +16,10 @@ def init_graph_store( import kuzu if db_path is None: - logging.info("No db_path provided, creating temporary directory for database") + logger.info("No db_path provided, creating temporary directory for database") db_path = os.path.join(tempfile.mkdtemp(), "kuzu_db") - logging.info("Using Kuzu graph store with path: %s", db_path) + logger.info("Using Kuzu graph store with path: %s", db_path) db = kuzu.Database(db_path) return MotleyKuzuGraphStore(db) diff --git a/motleycrew/storage/kuzu_graph_store.py b/motleycrew/storage/kuzu_graph_store.py index f2b04391..bb170af3 100644 --- a/motleycrew/storage/kuzu_graph_store.py +++ b/motleycrew/storage/kuzu_graph_store.py @@ -4,7 +4,6 @@ """ from typing import Any, Dict, List, Optional, Type, TypeVar -import logging from kuzu import Connection, PreparedStatement, QueryResult import json @@ -13,6 +12,7 @@ from motleycrew.storage import MotleyGraphStore from motleycrew.storage import MotleyGraphNode from motleycrew.storage import MotleyGraphNodeType +from motleycrew.common import logger class MotleyKuzuGraphStore(MotleyGraphStore): @@ -55,9 +55,9 @@ def _execute_query( """ Execute a query, logging it for debugging purposes """ - logging.debug("Executing query: %s", query) + logger.debug("Executing query: %s", query) if parameters: - logging.debug("with parameters: %s", parameters) + logger.debug("with parameters: %s", parameters) # TODO: retries? return self.connection.execute(query=query, parameters=parameters) @@ -91,7 +91,7 @@ def ensure_node_table(self, node_class: Type[MotleyGraphNode]) -> str: """ table_name = node_class.get_label() if not self._check_node_table_exists(table_name): - logging.info("Node table %s does not exist in the database, creating", table_name) + logger.info("Node table %s does not exist in the database, creating", table_name) self._execute_query( "CREATE NODE TABLE {} (id SERIAL, PRIMARY KEY(id))".format(table_name) ) @@ -100,7 +100,7 @@ def ensure_node_table(self, node_class: Type[MotleyGraphNode]) -> str: existing_property_names = self._get_node_property_names(node_class.get_label()) for field_name, field in node_class.model_fields.items(): if field_name not in existing_property_names: - logging.info( + logger.info( "Property %s not present in table for label %s, creating", field_name, node_class.get_label(), @@ -126,7 +126,7 @@ def ensure_relation_table( if not self._check_rel_table_exists( from_label=from_class.get_label(), to_label=to_class.get_label(), rel_label=label ): - logging.info( + logger.info( "Relation table %s from %s to %s does not exist in the database, creating", label, from_class.get_label(), @@ -229,7 +229,7 @@ def insert_node(self, node: MotleyGraphNodeType) -> MotleyGraphNodeType: assert node.id is None, "Entity has its id set, looks like it is already in the DB" self.ensure_node_table(type(node)) - logging.info("Inserting new node with label %s: %s", node.get_label(), node) + logger.info("Inserting new node with label %s: %s", node.get_label(), node) cypher_mapping, parameters = MotleyKuzuGraphStore._node_to_cypher_mapping_with_parameters( node @@ -239,7 +239,7 @@ def insert_node(self, node: MotleyGraphNodeType) -> MotleyGraphNodeType: parameters=parameters, ) assert create_result.has_next() - logging.info("Node created OK") + logger.info("Node created OK") created_object = create_result.get_next()[0] created_object_id = created_object.get("id") @@ -269,7 +269,7 @@ def create_relation( self.ensure_relation_table(from_class=type(from_node), to_class=type(to_node), label=label) - logging.info( + logger.info( "Creating relation %s from %s:%s to %s:%s", label, from_node.get_label(), @@ -290,7 +290,7 @@ def create_relation( }, ) assert create_result.has_next() - logging.info("Relation created OK") + logger.info("Relation created OK") def upsert_triplet(self, from_node: MotleyGraphNode, to_node: MotleyGraphNode, label: str): """ @@ -299,15 +299,15 @@ def upsert_triplet(self, from_node: MotleyGraphNode, to_node: MotleyGraphNode, l This method also creates and/or updates all necessary tables. """ if not self.check_node_exists(from_node): - logging.info("Node %s does not exist, creating", from_node) + logger.info("Node %s does not exist, creating", from_node) self.insert_node(from_node) if not self.check_node_exists(to_node): - logging.info("Node %s does not exist, creating", to_node) + logger.info("Node %s does not exist, creating", to_node) self.insert_node(to_node) if not self.check_relation_exists(from_node=from_node, to_node=to_node, label=label): - logging.info("Relation from %s to %s does not exist, creating", from_node, to_node) + logger.info("Relation from %s to %s does not exist, creating", from_node, to_node) self.create_relation(from_node=from_node, to_node=to_node, label=label) def delete_node(self, node: MotleyGraphNode) -> None: @@ -424,7 +424,7 @@ def _deserialize_node( if isinstance(value, str) and value.startswith( MotleyKuzuGraphStore.JSON_CONTENT_PREFIX ): - logging.debug( + logger.debug( "Value for field %s is marked as JSON, attempting to deserialize: %s", field_name, value, @@ -483,7 +483,7 @@ def _get_cypher_type_and_is_json_by_python_type_annotation( """ cypher_type = MotleyKuzuGraphStore.PYTHON_TO_CYPHER_TYPES_MAPPING.get(annotation) if not cypher_type: - logging.warning( + logger.warning( "No known Cypher type matching annotation %s, will use JSON string", annotation, ) diff --git a/motleycrew/tasks/simple.py b/motleycrew/tasks/simple.py index 8e56ea39..9746143d 100644 --- a/motleycrew/tasks/simple.py +++ b/motleycrew/tasks/simple.py @@ -1,5 +1,4 @@ from __future__ import annotations -import logging from typing import TYPE_CHECKING, Any, Sequence, List, Optional from motleycrew.tasks.task import Task @@ -7,6 +6,7 @@ from motleycrew.agents.abstract_parent import MotleyAgentAbstractParent from motleycrew.tools import MotleyTool +from motleycrew.common import logger if TYPE_CHECKING: from motleycrew.crew import MotleyCrew @@ -81,7 +81,7 @@ def register_completed_unit(self, unit: SimpleTaskUnit) -> None: def get_next_unit(self) -> SimpleTaskUnit | None: if self.done: - logging.info("Task %s is already done", self) + logger.info("Task %s is already done", self) return None upstream_tasks = self.get_upstream_tasks() @@ -102,12 +102,12 @@ def get_worker(self, tools: Optional[List[MotleyTool]]) -> MotleyAgentAbstractPa raise ValueError("Task is not associated with an agent") if hasattr(self.agent, "is_materialized") and self.agent.is_materialized and tools: - logging.warning( + logger.warning( "Agent %s is already materialized, can't add extra tools %s", self.agent, tools ) if hasattr(self.agent, "add_tools") and tools: - logging.info("Adding tools %s to agent %s", tools, self.agent) + logger.info("Adding tools %s to agent %s", tools, self.agent) self.agent.add_tools(tools) # TODO: that's a pretty big assumption on agent structure. Necessary? diff --git a/motleycrew/tools/image/dall_e.py b/motleycrew/tools/image/dall_e.py index 951c43c7..2c967908 100644 --- a/motleycrew/tools/image/dall_e.py +++ b/motleycrew/tools/image/dall_e.py @@ -1,4 +1,3 @@ -import logging from typing import Optional import os @@ -13,6 +12,7 @@ import motleycrew.common.utils as motley_utils from motleycrew.common import LLMFramework from motleycrew.common.llms import init_llm +from motleycrew.common import logger from langchain.prompts import PromptTemplate @@ -25,7 +25,7 @@ def download_image(url: str, file_path: str) -> Optional[str]: extension = ".png" # default to .png if content-type is not recognized file_path_with_extension = file_path + extension - logging.info("Downloading image %s to %s", url, file_path_with_extension) + logger.info("Downloading image %s to %s", url, file_path_with_extension) with open(file_path_with_extension, "wb") as f: for chunk in response: @@ -33,7 +33,7 @@ def download_image(url: str, file_path: str) -> Optional[str]: return file_path_with_extension else: - logging.error("Failed to download image. Status code: %s", response.status_code) + logger.error("Failed to download image. Status code: %s", response.status_code) class DallEImageGeneratorTool(MotleyTool): @@ -101,11 +101,11 @@ def run_dalle_and_save_images( ) dalle_result = dalle_api.run(prompt_value.text) - logging.info("Dall-E API output: %s", dalle_result) + logger.info("Dall-E API output: %s", dalle_result) urls = dalle_result.split(dalle_api.separator) if not len(urls) or not motley_utils.is_http_url(urls[0]): - logging.error("Dall-E API did not return a valid url: %s", dalle_result) + logger.error("Dall-E API did not return a valid url: %s", dalle_result) return if images_directory: @@ -121,7 +121,7 @@ def run_dalle_and_save_images( file_paths.append(file_path_with_extension) return file_paths else: - logging.info("Images directory is not provided, returning URLs") + logger.info("Images directory is not provided, returning URLs") return urls @@ -156,4 +156,4 @@ def run_dalle_and_save_images_partial(description: str): if __name__ == "__main__": tool = DallEImageGeneratorTool() out = tool.invoke("A beautiful castle on top of a hill at sunset") - logging.info(out) + logger.info(out) diff --git a/motleycrew/tracking/callbacks.py b/motleycrew/tracking/callbacks.py index decafe6d..676c0c45 100644 --- a/motleycrew/tracking/callbacks.py +++ b/motleycrew/tracking/callbacks.py @@ -1,5 +1,4 @@ from typing import List, Dict, Optional, Any, Union -import logging import traceback try: @@ -16,6 +15,7 @@ from motleycrew.common.enums import LunaryRunType, LunaryEventName from motleycrew.common.utils import ensure_module_is_installed +from motleycrew.common import logger def event_delegate_decorator(f): @@ -47,12 +47,12 @@ def wrapper(self, *args, **kwargs): msg = "[Lunary] An error occurred in {}: {}\n{}".format( handler_name, e, traceback.format_exc() ) - logging.warning(msg) + logger.warning(msg) else: msg = "No handler with the name of the {} was found for {}".format( handler_name, self.__class__ ) - logging.warning(msg) + logger.warning(msg) return f(*args, **kwargs) diff --git a/motleycrew/tracking/utils.py b/motleycrew/tracking/utils.py index 7e866ceb..ba45638c 100644 --- a/motleycrew/tracking/utils.py +++ b/motleycrew/tracking/utils.py @@ -1,6 +1,5 @@ import os from typing import List, Optional -import logging from lunary import LunaryCallbackHandler from langchain_core.callbacks import BaseCallbackHandler, BaseCallbackManager @@ -8,6 +7,7 @@ from .callbacks import LlamaIndexLunaryCallbackHandler from motleycrew.common import LLMFramework +from motleycrew.common import logger def get_lunary_public_key(): @@ -18,7 +18,7 @@ def get_lunary_public_key(): or os.getenv("LLMONITOR_APP_ID") ) if not key: - logging.warning("Lunary public key is not set, tracking will be disabled") + logger.warning("Lunary public key is not set, tracking will be disabled") return key @@ -70,7 +70,7 @@ def get_default_callbacks_list( _default_callbacks = dc_factory() else: msg = "Default callbacks are not implemented for {} framework".format(framework_name) - logging.warning(msg) + logger.warning(msg) return _default_callbacks diff --git a/tests/run_integration_tests.py b/tests/run_integration_tests.py index ca9ddb37..ac2938de 100644 --- a/tests/run_integration_tests.py +++ b/tests/run_integration_tests.py @@ -4,7 +4,6 @@ import os import argparse from pathlib import Path -import logging import traceback import difflib import json @@ -16,7 +15,7 @@ from nbconvert.preprocessors import ExecutePreprocessor from motleycrew.common.exceptions import IntegrationTestException -from motleycrew.common.utils import configure_logging +from motleycrew.common import logger, configure_logging from motleycache import ( enable_cache, @@ -158,7 +157,7 @@ def build_ipynb_integration_tests(is_minimal: bool = False) -> dict: tests = MINIMAL_IPYNB_INTEGRATION_TESTS if is_minimal else IPYNB_INTEGRATION_TESTS for test_name, nb_path in tests.items(): if not os.path.exists(nb_path): - logging.info("Ipynb test notebook {} not found".format(test_name)) + logger.info("Ipynb test notebook {} not found".format(test_name)) continue test_functions[test_name] = partial(run_ipynb, nb_path) @@ -194,11 +193,11 @@ def run_integration_tests( if test_name is not None and test_name != current_test_name: continue - logging.info("Running test: %s", current_test_name) + logger.info("Running test: %s", current_test_name) cache_sub_dir = os.path.join(cache_dir, current_test_name) if update_golden: - logging.info("Update-golden flag is set. Cleaning cache directory %s", cache_sub_dir) + logger.info("Update-golden flag is set. Cleaning cache directory %s", cache_sub_dir) shutil.rmtree(cache_sub_dir, ignore_errors=True) os.makedirs(cache_sub_dir, exist_ok=True) os.makedirs(golden_dir, exist_ok=True) @@ -211,7 +210,7 @@ def run_integration_tests( test_result = test_fn() if current_test_name in INTEGRATION_TESTS: if update_golden: - logging.info( + logger.info( "Skipping check and updating golden data for test: %s", current_test_name ) write_content(golden_dir, current_test_name, test_result) @@ -220,17 +219,17 @@ def run_integration_tests( compare_results(test_result, excepted_result) except Exception as e: - logging.error("Test %s failed: %s", current_test_name, str(e)) + logger.error("Test %s failed: %s", current_test_name, str(e)) failed_tests[current_test_name] = traceback.format_exc() for t, exception in failed_tests.items(): - logging.error("Test %s failed", t) - logging.error(exception) + logger.error("Test %s failed", t) + logger.error(exception) if failed_tests: raise IntegrationTestException(test_names=list(failed_tests.keys())) - logging.info("All tests passed!") + logger.info("All tests passed!") unset_tiktoken_cache_dir(old_tiktoken_cache_dir)