Skip to content

Commit 6340e87

Browse files
hassiebphassiebp
committed
push
1 parent a93acca commit 6340e87

1 file changed

Lines changed: 73 additions & 36 deletions

File tree

‎tests/test_langchain_integration.py‎

Lines changed: 73 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from langchain.schema import StrOutputParser
66
from langchain_openai import ChatOpenAI, OpenAI
77

8+
from langfuse import Langfuse
89
from langfuse.langchain import CallbackHandler
910
from tests.utils import get_api
1011

@@ -28,8 +29,10 @@ def test_stream_chat_models(model_name):
2829
handler = CallbackHandler()
2930

3031
langfuse_client = handler.client
31-
with langfuse_client.start_as_current_span(name=name) as span:
32-
trace_id = span.trace_id
32+
trace_id = Langfuse.create_trace_id()
33+
with langfuse_client.start_as_current_span(
34+
name=name, trace_context={"trace_id": trace_id}
35+
):
3336
res = model.stream(
3437
[{"role": "user", "content": "return the exact phrase - This is a test!"}],
3538
config={"callbacks": [handler]},
@@ -78,8 +81,10 @@ def test_stream_completions_models(model_name):
7881
handler = CallbackHandler()
7982

8083
langfuse_client = handler.client
81-
with langfuse_client.start_as_current_span(name=name) as span:
82-
trace_id = span.trace_id
84+
trace_id = Langfuse.create_trace_id()
85+
with langfuse_client.start_as_current_span(
86+
name=name, trace_context={"trace_id": trace_id}
87+
):
8388
res = model.stream(
8489
"return the exact phrase - This is a test!",
8590
config={"callbacks": [handler]},
@@ -127,8 +132,10 @@ def test_invoke_chat_models(model_name):
127132
handler = CallbackHandler()
128133

129134
langfuse_client = handler.client
130-
with langfuse_client.start_as_current_span(name=name) as span:
131-
trace_id = span.trace_id
135+
trace_id = Langfuse.create_trace_id()
136+
with langfuse_client.start_as_current_span(
137+
name=name, trace_context={"trace_id": trace_id}
138+
):
132139
_ = model.invoke(
133140
[{"role": "user", "content": "return the exact phrase - This is a test!"}],
134141
config={"callbacks": [handler]},
@@ -172,8 +179,10 @@ def test_invoke_in_completions_models(model_name):
172179
handler = CallbackHandler()
173180

174181
langfuse_client = handler.client
175-
with langfuse_client.start_as_current_span(name=name) as span:
176-
trace_id = span.trace_id
182+
trace_id = Langfuse.create_trace_id()
183+
with langfuse_client.start_as_current_span(
184+
name=name, trace_context={"trace_id": trace_id}
185+
):
177186
test_phrase = "This is a test!"
178187
_ = model.invoke(
179188
f"return the exact phrase - {test_phrase}",
@@ -216,8 +225,10 @@ def test_batch_in_completions_models(model_name):
216225
handler = CallbackHandler()
217226

218227
langfuse_client = handler.client
219-
with langfuse_client.start_as_current_span(name=name) as span:
220-
trace_id = span.trace_id
228+
trace_id = Langfuse.create_trace_id()
229+
with langfuse_client.start_as_current_span(
230+
name=name, trace_context={"trace_id": trace_id}
231+
):
221232
input1 = "Who is the first president of America ?"
222233
input2 = "Who is the first president of Ireland ?"
223234
_ = model.batch(
@@ -260,8 +271,10 @@ def test_batch_in_chat_models(model_name):
260271
handler = CallbackHandler()
261272

262273
langfuse_client = handler.client
263-
with langfuse_client.start_as_current_span(name=name) as span:
264-
trace_id = span.trace_id
274+
trace_id = Langfuse.create_trace_id()
275+
with langfuse_client.start_as_current_span(
276+
name=name, trace_context={"trace_id": trace_id}
277+
):
265278
input1 = "Who is the first president of America ?"
266279
input2 = "Who is the first president of Ireland ?"
267280
_ = model.batch(
@@ -307,8 +320,10 @@ async def test_astream_chat_models(model_name):
307320
handler = CallbackHandler()
308321

309322
langfuse_client = handler.client
310-
with langfuse_client.start_as_current_span(name=name) as span:
311-
trace_id = span.trace_id
323+
trace_id = Langfuse.create_trace_id()
324+
with langfuse_client.start_as_current_span(
325+
name=name, trace_context={"trace_id": trace_id}
326+
):
312327
res = model.astream(
313328
[{"role": "user", "content": "Who was the first American president "}],
314329
config={"callbacks": [handler]},
@@ -358,8 +373,10 @@ async def test_astream_completions_models(model_name):
358373

359374
langfuse_client = handler.client
360375

361-
with langfuse_client.start_as_current_span(name=name) as span:
362-
trace_id = span.trace_id
376+
trace_id = Langfuse.create_trace_id()
377+
with langfuse_client.start_as_current_span(
378+
name=name, trace_context={"trace_id": trace_id}
379+
):
363380
test_phrase = "This is a test!"
364381
res = model.astream(
365382
f"return the exact phrase - {test_phrase}",
@@ -409,8 +426,10 @@ async def test_ainvoke_chat_models(model_name):
409426
handler = CallbackHandler()
410427

411428
langfuse_client = handler.client
412-
with langfuse_client.start_as_current_span(name=name) as span:
413-
trace_id = span.trace_id
429+
trace_id = Langfuse.create_trace_id()
430+
with langfuse_client.start_as_current_span(
431+
name=name, trace_context={"trace_id": trace_id}
432+
):
414433
test_phrase = "This is a test!"
415434
_ = await model.ainvoke(
416435
[{"role": "user", "content": f"return the exact phrase - {test_phrase} "}],
@@ -455,8 +474,10 @@ async def test_ainvoke_in_completions_models(model_name):
455474
handler = CallbackHandler()
456475

457476
langfuse_client = handler.client
458-
with langfuse_client.start_as_current_span(name=name) as span:
459-
trace_id = span.trace_id
477+
trace_id = Langfuse.create_trace_id()
478+
with langfuse_client.start_as_current_span(
479+
name=name, trace_context={"trace_id": trace_id}
480+
):
460481
test_phrase = "This is a test!"
461482
_ = await model.ainvoke(
462483
f"return the exact phrase - {test_phrase}",
@@ -503,8 +524,10 @@ def test_chains_batch_in_chat_models(model_name):
503524
handler = CallbackHandler()
504525

505526
langfuse_client = handler.client
506-
with langfuse_client.start_as_current_span(name=name) as span:
507-
trace_id = span.trace_id
527+
trace_id = Langfuse.create_trace_id()
528+
with langfuse_client.start_as_current_span(
529+
name=name, trace_context={"trace_id": trace_id}
530+
):
508531
prompt = ChatPromptTemplate.from_template(
509532
"tell me a joke about {foo} in 300 words"
510533
)
@@ -549,8 +572,10 @@ def test_chains_batch_in_completions_models(model_name):
549572
handler = CallbackHandler()
550573

551574
langfuse_client = handler.client
552-
with langfuse_client.start_as_current_span(name=name) as span:
553-
trace_id = span.trace_id
575+
trace_id = Langfuse.create_trace_id()
576+
with langfuse_client.start_as_current_span(
577+
name=name, trace_context={"trace_id": trace_id}
578+
):
554579
prompt = ChatPromptTemplate.from_template(
555580
"tell me a joke about {foo} in 300 words"
556581
)
@@ -597,8 +622,10 @@ async def test_chains_abatch_in_chat_models(model_name):
597622
handler = CallbackHandler()
598623

599624
langfuse_client = handler.client
600-
with langfuse_client.start_as_current_span(name=name) as span:
601-
trace_id = span.trace_id
625+
trace_id = Langfuse.create_trace_id()
626+
with langfuse_client.start_as_current_span(
627+
name=name, trace_context={"trace_id": trace_id}
628+
):
602629
prompt = ChatPromptTemplate.from_template(
603630
"tell me a joke about {foo} in 300 words"
604631
)
@@ -645,8 +672,10 @@ async def test_chains_abatch_in_completions_models(model_name):
645672
handler = CallbackHandler()
646673

647674
langfuse_client = handler.client
648-
with langfuse_client.start_as_current_span(name=name) as span:
649-
trace_id = span.trace_id
675+
trace_id = Langfuse.create_trace_id()
676+
with langfuse_client.start_as_current_span(
677+
name=name, trace_context={"trace_id": trace_id}
678+
):
650679
prompt = ChatPromptTemplate.from_template(
651680
"tell me a joke about {foo} in 300 words"
652681
)
@@ -689,8 +718,10 @@ async def test_chains_ainvoke_chat_models(model_name):
689718
handler = CallbackHandler()
690719

691720
langfuse_client = handler.client
692-
with langfuse_client.start_as_current_span(name=name) as span:
693-
trace_id = span.trace_id
721+
trace_id = Langfuse.create_trace_id()
722+
with langfuse_client.start_as_current_span(
723+
name=name, trace_context={"trace_id": trace_id}
724+
):
694725
prompt1 = ChatPromptTemplate.from_template(
695726
"""You are a skilled writer tasked with crafting an engaging introduction for a blog post on the following topic:
696727
Topic: {topic}
@@ -740,8 +771,10 @@ async def test_chains_ainvoke_completions_models(model_name):
740771
handler = CallbackHandler()
741772

742773
langfuse_client = handler.client
743-
with langfuse_client.start_as_current_span(name=name) as span:
744-
trace_id = span.trace_id
774+
trace_id = Langfuse.create_trace_id()
775+
with langfuse_client.start_as_current_span(
776+
name=name, trace_context={"trace_id": trace_id}
777+
):
745778
prompt1 = PromptTemplate.from_template(
746779
"""You are a skilled writer tasked with crafting an engaging introduction for a blog post on the following topic:
747780
Topic: {topic}
@@ -791,8 +824,10 @@ async def test_chains_astream_chat_models(model_name):
791824
handler = CallbackHandler()
792825

793826
langfuse_client = handler.client
794-
with langfuse_client.start_as_current_span(name=name) as span:
795-
trace_id = span.trace_id
827+
trace_id = Langfuse.create_trace_id()
828+
with langfuse_client.start_as_current_span(
829+
name=name, trace_context={"trace_id": trace_id}
830+
):
796831
prompt1 = PromptTemplate.from_template(
797832
"""You are a skilled writer tasked with crafting an engaging introduction for a blog post on the following topic:
798833
Topic: {topic}
@@ -848,8 +883,10 @@ async def test_chains_astream_completions_models(model_name):
848883
handler = CallbackHandler()
849884

850885
langfuse_client = handler.client
851-
with langfuse_client.start_as_current_span(name=name) as span:
852-
trace_id = span.trace_id
886+
trace_id = Langfuse.create_trace_id()
887+
with langfuse_client.start_as_current_span(
888+
name=name, trace_context={"trace_id": trace_id}
889+
):
853890
prompt1 = PromptTemplate.from_template(
854891
"""You are a skilled writer tasked with crafting an engaging introduction for a blog post on the following topic:
855892
Topic: {topic}

0 commit comments

Comments
 (0)