removed more than 15 local vars from testing file and changed return logic for init file

Author: haneric00

ai_agent_instrumentation/opentelemetry-instrumentation-langchain-v2/src/opentelemetry/instrumentation/langchain_v2/__init__.py

@@ -60,3 +60,4 @@ def __call__(
                 return None
 
         instance.add_handler(self.callback_handler, True)
+        return None

ai_agent_instrumentation/opentelemetry-instrumentation-langchain-v2/tests/test_chains.py

@@ -14,48 +14,68 @@
 from opentelemetry.trace import SpanKind
 
 
-@pytest.mark.vcr(filter_headers=["Authorization", "X-Amz-Date", "X-Amz-Security-Token"], record_mode="once")
-def test_sequential_chain(instrument_langchain, span_exporter):
-    span_exporter.clear()
-
-    session = boto3.Session(region_name="us-west-2")
-
-    bedrock_client = session.client(service_name="bedrock-runtime", region_name="us-west-2")
-
-    llm = BedrockLLM(
+def create_bedrock_llm(region="us-west-2"):
+    """Create and return a BedrockLLM instance."""
+    session = boto3.Session(region_name=region)
+    bedrock_client = session.client(service_name="bedrock-runtime", region_name=region)
+    return BedrockLLM(
         client=bedrock_client,
         model_id="anthropic.claude-v2",
-        model_kwargs={
-            "max_tokens_to_sample": 500,
-            "temperature": 0.7,
-        },
+        model_kwargs={"max_tokens_to_sample": 500, "temperature": 0.7},
     )
-    synopsis_template = """You are a playwright. Given the title of play and the era it is set in, it is your job to write a synopsis for that title.
+
+
+def create_chains(llm):
+    """Create and return the synopsis chain, review chain, and overall chain."""
+
+    synopsis_prompt = PromptTemplate(
+        input_variables=["title", "era"],
+        template="""You are a playwright. Given the title of play and the era it is set in, it is your job to write a synopsis for that title.
 
     Title: {title}
     Era: {era}
-    Playwright: This is a synopsis for the above play:"""  # noqa: E501
-    synopsis_prompt_template = PromptTemplate(input_variables=["title", "era"], template=synopsis_template)
-    synopsis_chain = LLMChain(llm=llm, prompt=synopsis_prompt_template, output_key="synopsis", name="synopsis")
+    Playwright: This is a synopsis for the above play:""",  # noqa: E501
+    )
+    synopsis_chain = LLMChain(llm=llm, prompt=synopsis_prompt, output_key="synopsis", name="synopsis")
 
-    template = """You are a play critic from the New York Times. Given the synopsis of play, it is your job to write a review for that play.
+    review_prompt = PromptTemplate(
+        input_variables=["synopsis"],
+        template="""You are a play critic from the New York Times. Given the synopsis of play, it is your job to write a review for that play.
 
     Play Synopsis:
     {synopsis}
-    Review from a New York Times play critic of the above play:"""  # noqa: E501
-    prompt_template = PromptTemplate(input_variables=["synopsis"], template=template)
-    review_chain = LLMChain(llm=llm, prompt=prompt_template, output_key="review")
+    Review from a New York Times play critic of the above play:""",  # noqa: E501
+    )
+    review_chain = LLMChain(llm=llm, prompt=review_prompt, output_key="review")
 
     overall_chain = SequentialChain(
         chains=[synopsis_chain, review_chain],
         input_variables=["era", "title"],
         output_variables=["synopsis", "review"],
         verbose=True,
     )
-    overall_chain.invoke({"title": "Tragedy at sunset on the beach", "era": "Victorian England"})
 
-    spans = span_exporter.get_finished_spans()
+    return overall_chain
+
+
+def validate_span(span, expected_kind, expected_attrs):
+    """Validate a span against expected values."""
+    assert span.kind == expected_kind
+    for attr in expected_attrs:
+        assert attr in span.attributes
+    return ast.literal_eval(span.attributes["gen_ai.prompt"]), ast.literal_eval(span.attributes["gen_ai.completion"])
+
 
+@pytest.mark.vcr(filter_headers=["Authorization", "X-Amz-Date", "X-Amz-Security-Token"], record_mode="once")
+def test_sequential_chain(instrument_langchain, span_exporter):
+    span_exporter.clear()
+
+    llm = create_bedrock_llm()
+    chain = create_chains(llm)
+    input_data = {"title": "Tragedy at sunset on the beach", "era": "Victorian England"}
+    chain.invoke(input_data)
+
+    spans = span_exporter.get_finished_spans()
     langchain_spans = [span for span in spans if span.name.startswith("chain ")]
 
     assert [
@@ -68,36 +88,23 @@ def test_sequential_chain(instrument_langchain, span_exporter):
     review_span = next(span for span in spans if span.name == "chain LLMChain")
     overall_span = next(span for span in spans if span.name == "chain SequentialChain")
 
-    assert synopsis_span.kind == SpanKind.INTERNAL
-    assert "gen_ai.prompt" in synopsis_span.attributes
-    assert "gen_ai.completion" in synopsis_span.attributes
-
-    synopsis_prompt = ast.literal_eval(synopsis_span.attributes["gen_ai.prompt"])
-    synopsis_completion = ast.literal_eval(synopsis_span.attributes["gen_ai.completion"])
-
-    assert synopsis_prompt == {"title": "Tragedy at sunset on the beach", "era": "Victorian England"}
+    synopsis_prompt, synopsis_completion = validate_span(
+        synopsis_span, SpanKind.INTERNAL, ["gen_ai.prompt", "gen_ai.completion"]
+    )
+    assert synopsis_prompt == input_data
     assert "synopsis" in synopsis_completion
 
-    assert review_span.kind == SpanKind.INTERNAL
-    assert "gen_ai.prompt" in review_span.attributes
-    assert "gen_ai.completion" in review_span.attributes
-    print("Raw completion value:", repr(synopsis_span.attributes["gen_ai.completion"]))
-
-    review_prompt = ast.literal_eval(review_span.attributes["gen_ai.prompt"])
-    review_completion = ast.literal_eval(review_span.attributes["gen_ai.completion"])
-
+    review_prompt, review_completion = validate_span(
+        review_span, SpanKind.INTERNAL, ["gen_ai.prompt", "gen_ai.completion"]
+    )
     assert "title" in review_prompt
     assert "era" in review_prompt
     assert "synopsis" in review_prompt
     assert "review" in review_completion
 
-    assert overall_span.kind == SpanKind.INTERNAL
-    assert "gen_ai.prompt" in overall_span.attributes
-    assert "gen_ai.completion" in overall_span.attributes
-
-    overall_prompt = ast.literal_eval(overall_span.attributes["gen_ai.prompt"])
-    overall_completion = ast.literal_eval(overall_span.attributes["gen_ai.completion"])
-
-    assert overall_prompt == {"title": "Tragedy at sunset on the beach", "era": "Victorian England"}
+    overall_prompt, overall_completion = validate_span(
+        overall_span, SpanKind.INTERNAL, ["gen_ai.prompt", "gen_ai.completion"]
+    )
+    assert overall_prompt == input_data
     assert "synopsis" in overall_completion
     assert "review" in overall_completion

ai_agent_instrumentation/opentelemetry-instrumentation-langchain-v2/tests/test_langgraph_agent.py

@@ -56,9 +56,9 @@ def calculate(state: State):
         print(f"  Attributes: {span.attributes}")
         print("---")
 
-    expected_spans = {"chain LangGraph", "chain calculate", "chat anthropic.claude-3-haiku-20240307-v1:0"}
-
-    assert expected_spans == {span.name for span in spans}
+    assert {"chain LangGraph", "chain calculate", "chat anthropic.claude-3-haiku-20240307-v1:0"} == {
+        span.name for span in spans
+    }
 
     llm_span = next(span for span in spans if span.name == "chat anthropic.claude-3-haiku-20240307-v1:0")
     calculate_task_span = next(span for span in spans if span.name == "chain calculate")
@@ -115,8 +115,7 @@ def calculate(state: State):
 
     langgraph = workflow.compile()
 
-    user_request = "What's 5 + 5?"
-    await langgraph.ainvoke(input={"request": user_request})
+    await langgraph.ainvoke(input={"request": "What's 5 + 5?"})
     spans = span_exporter.get_finished_spans()
 
     assert set(["chain LangGraph", "chain calculate", "chat anthropic.claude-3-haiku-20240307-v1:0"]) == {
@@ -147,8 +146,6 @@ def build_graph():
 
     graph = build_graph()
 
-    from opentelemetry import trace
-
     assert trace.get_current_span() == INVALID_SPAN
 
     graph.invoke({"result": "init"})