add MLflow operator CLI

Author: ks6088ts

Makefile

@@ -173,7 +173,10 @@ create-elasticsearch-index: ## create Elasticsearch index
 # ---
 # Project / Run agents
 # ---
+
 QUESTION ?= "KABUTOの起動時に、画面全体が紫色に点滅し、システムがフリーズします。KABUTO のマニュアルから、関連する情報を取得したり過去のシステムのトラブルシュート事例が蓄積されたデータベースから、関連する情報を取得して質問に答えてください"
+PORT ?= 5001
+
 .PHONY: run-chat-with-tools-agent
 run-chat-with-tools-agent: ## run chat with tools agent
 	uv run python scripts/agent_operator.py run \
@@ -227,4 +230,4 @@ mlflow: ## run MLflow
 	uv run mlflow server \
 		--backend-store-uri sqlite:///mlflow.db \
 		--host 0.0.0.0 \
-		--port 5000
+		--port $(PORT)

docs/references.md

@@ -90,4 +90,9 @@
 - [DSPy (Declarative Self-improving Python)](https://dspy.ai/)
 - [Language Models](https://dspy.ai/learn/programming/language_models/)
 - [Language Models / v3.0.3](https://github.com/stanfordnlp/dspy/blob/3.0.3/docs/docs/learn/programming/language_models.md)
-- [Software Design誌「実践LLMアプリケーション開発」第25回サンプルコード](https://github.com/mahm/softwaredesign-llm-application/tree/main/25)
+- [Software Design 誌「実践 LLM アプリケーション開発」第 25 回サンプルコード](https://github.com/mahm/softwaredesign-llm-application/tree/main/25)
+
+### [MLflow](https://mlflow.org/docs/latest/genai/)
+
+- [LangChain / MLflow](https://docs.langchain.com/oss/python/integrations/providers/mlflow_tracking)
+- [MLflow / Tracing LangGraph🦜🕸️](https://mlflow.org/docs/latest/genai/tracing/integrations/listing/langgraph/)

scripts/dspy_operator.py

@@ -76,7 +76,7 @@ def optimize_with_miprov2(trainset, eval_lm, chat_lm):
     """MIPROv2を使用してチャットボットを最適化"""
 
     # MLflowの設定
-    MLFLOW_PORT = os.getenv("MLFLOW_PORT", "5000")
+    MLFLOW_PORT = os.getenv("MLFLOW_PORT", "5001")
     MLFLOW_TRACKING_URI = f"http://localhost:{MLFLOW_PORT}"
     MLFLOW_EXPERIMENT_NAME = "DSPy-EdamameFairy-Optimization"
     MLFLOW_RUN_NAME = "miprov2_optimization"

scripts/mlflow_operator.py

@@ -0,0 +1,76 @@
+import logging
+from logging import basicConfig
+
+import mlflow
+import typer
+from dotenv import load_dotenv
+from langchain_core.messages import HumanMessage
+
+from template_langgraph.agents.demo_agents.weather_agent import graph
+from template_langgraph.loggers import get_logger
+
+app = typer.Typer(
+    add_completion=False,
+    help="MLflow operator CLI",
+)
+logger = get_logger(__name__)
+
+
+def set_verbose_logging(verbose: bool):
+    if verbose:
+        logger.setLevel(logging.DEBUG)
+        basicConfig(level=logging.DEBUG)
+
+
+@app.command(
+    help="Run the LangGraph agent with MLflow tracing ref. https://mlflow.org/docs/2.21.3/tracing/integrations/langgraph"
+)
+def tracing(
+    query: str = typer.Option(
+        "What is the weather like in Japan?",
+        "--query",
+        "-q",
+        help="Query to run with the LangGraph agent",
+    ),
+    tracking_uri: str = typer.Option(
+        "http://localhost:5001",
+        "--tracking-uri",
+        "-t",
+        help="MLflow tracking URI",
+    ),
+    verbose: bool = typer.Option(
+        True,
+        "--verbose",
+        "-v",
+        help="Enable verbose output",
+    ),
+):
+    set_verbose_logging(verbose)
+    logger.info("Running...")
+
+    mlflow.langchain.autolog()
+    mlflow.set_tracking_uri(tracking_uri)
+    mlflow.set_experiment("LangGraph Experiment")
+
+    result = graph.invoke(
+        {
+            "messages": [
+                HumanMessage(content=query),
+            ]
+        },
+    )
+    logger.info(f"Result: {result}")
+
+    # Get the trace object just created
+    trace = mlflow.get_trace(
+        trace_id=mlflow.get_last_active_trace_id(),
+    )
+    logger.info(f"Trace info: {trace.info.token_usage}")
+
+
+if __name__ == "__main__":
+    load_dotenv(
+        override=True,
+        verbose=True,
+    )
+    app()