fix(openai): langchain streaming bug (#225)

packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/__init__.py

@@ -205,6 +205,7 @@ def _set_response_attributes(span, llm_request_type, response):
 def _build_from_streaming_response(span, llm_request_type, response):
     complete_response = {"choices": [], "model": ""}
     for item in response:
+        item_to_yield = item
         if is_openai_v1():
             item = item.__dict__
 
@@ -234,7 +235,7 @@ def _build_from_streaming_response(span, llm_request_type, response):
             else:
                 complete_choice["text"] += choice.get("text")
 
-        yield item
+        yield item_to_yield
 
     _set_response_attributes(
         span,

packages/traceloop-sdk/tests/test_langchain_instrumentation.py

@@ -1,5 +1,8 @@
 from langchain.prompts import PromptTemplate
 from langchain.llms import OpenAI
+from langchain.schema import StrOutputParser
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import PromptTemplate
 from langchain import LLMChain
 from langchain.chains import SequentialChain
 
@@ -38,7 +41,6 @@ def test_langchain(exporter):
     )
 
     spans = exporter.get_finished_spans()
-    print([span.name for span in spans])
 
     assert set(
         [
@@ -48,3 +50,25 @@ def test_langchain(exporter):
             "langchain.workflow",
         ]
     ).issubset([span.name for span in spans])
+
+
+def test_langchain_streaming(exporter):
+    chat = ChatOpenAI(
+        model="gpt-4",
+        temperature=0,
+        streaming=True,
+    )
+
+    prompt = PromptTemplate.from_template(
+        "write 10 lines of random text about ${product}"
+    )
+    runnable = prompt | chat | StrOutputParser()
+    runnable.invoke({"product": "colorful socks"})
+
+    spans = exporter.get_finished_spans()
+
+    assert set(
+        [
+            "openai.chat",
+        ]
+    ).issubset([span.name for span in spans])