rename callback.client

Author: hassiebp

langfuse/langchain/CallbackHandler.py

@@ -57,7 +57,7 @@
 
 class LangchainCallbackHandler(LangchainBaseCallbackHandler):
     def __init__(self, *, public_key: Optional[str] = None) -> None:
-        self.langfuse_client = get_client(public_key=public_key)
+        self.client = get_client(public_key=public_key)
 
         self.runs: Dict[UUID, Union[LangfuseSpan, LangfuseGeneration]] = {}
         self.prompt_to_parent_run_map = {}
@@ -174,7 +174,7 @@ def on_chain_start(
             }
 
             if parent_run_id is None:
-                self.runs[run_id] = self.langfuse_client.start_span(**content)
+                self.runs[run_id] = self.client.start_span(**content)
             else:
                 self.runs[run_id] = cast(
                     LangfuseSpan, self.runs[parent_run_id]
@@ -439,7 +439,7 @@ def on_retriever_start(
                     "level": "DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None,
                 }
 
-                self.runs[run_id] = self.langfuse_client.start_span(**content)
+                self.runs[run_id] = self.client.start_span(**content)
             else:
                 self.runs[run_id] = cast(
                     LangfuseSpan, self.runs[parent_run_id]
@@ -562,7 +562,7 @@ def __on_llm_action(
                     LangfuseSpan, self.runs[parent_run_id]
                 ).start_generation(**content)
             else:
-                self.runs[run_id] = self.langfuse_client.start_generation(**content)
+                self.runs[run_id] = self.client.start_generation(**content)
 
         except Exception as e:
             langfuse_logger.exception(e)

tests/test_extract_model.py

@@ -135,7 +135,7 @@ def test_models(expected_model: str, model: Any):
 def test_entire_llm_call(expected_model, model):
     callback = CallbackHandler()
 
-    with callback.langfuse_client.start_as_current_span(name="parent") as span:
+    with callback.client.start_as_current_span(name="parent") as span:
         trace_id = span.trace_id
 
         try:
@@ -146,7 +146,7 @@ def test_entire_llm_call(expected_model, model):
             print(e)
             pass
 
-    callback.langfuse_client.flush()
+    callback.client.flush()
     api = get_api()
 
     trace = api.trace.get(trace_id)

tests/test_extract_model_langchain_openai.py

@@ -327,7 +327,7 @@ def test_callback_retriever_conversational_with_memory():
         )
         conversation.predict(input="Hi there!", callbacks=[handler])
 
-    handler.langfuse_client.flush()
+    handler.client.flush()
 
     trace = get_api().trace.get(trace_id)
 
@@ -378,7 +378,7 @@ def test_callback_retriever_conversational():
 
         chain({"question": query, "chat_history": []}, callbacks=[handler])
 
-    handler.langfuse_client.flush()
+    handler.client.flush()
 
     trace = api_wrapper.get_trace(trace_id)
 
@@ -415,7 +415,7 @@ def test_callback_simple_openai():
         llm.invoke(text, config={"callbacks": [handler], "run_name": test_name})
 
         # Ensure data is flushed to API
-    handler.langfuse_client.flush()
+    handler.client.flush()
     sleep(2)
 
     # Retrieve trace
@@ -459,7 +459,7 @@ def test_callback_multiple_invocations_on_different_traces():
         handler2 = CallbackHandler()
         llm.invoke(text, config={"callbacks": [handler2], "run_name": test_name_2})
 
-    handler1.langfuse_client.flush()
+    handler1.client.flush()
 
     # Ensure data is flushed to API
     sleep(2)
@@ -543,7 +543,7 @@ def record_dog(name: str, color: str, fav_food: OptionalFavFood) -> str:
             callbacks=[handler],
         )
 
-    handler.langfuse_client.flush()
+    handler.client.flush()
 
     trace = get_api().trace.get(trace_id)
 
@@ -635,7 +635,7 @@ def get_word_length(word: str) -> int:
             config={"callbacks": [callback]},
         )
 
-    callback.langfuse_client.flush()
+    callback.client.flush()
 
     trace = get_api().trace.get(trace_id)
 
@@ -710,7 +710,7 @@ def _identifying_params(self) -> Mapping[str, Any]:
         )
         sequential_chain.run("This is a foobar thing", callbacks=[callback])
 
-    callback.langfuse_client.flush()
+    callback.client.flush()
 
     trace = get_api().trace.get(trace_id)
 
@@ -763,7 +763,7 @@ def test_openai_instruct_usage():
         ]
         runnable_chain.batch(input_list)
 
-    lf_handler.langfuse_client.flush()
+    lf_handler.client.flush()
 
     observations = get_api().trace.get(trace_id).observations
 
@@ -933,7 +933,7 @@ def test_link_langfuse_prompts_invoke():
             },
         )
 
-    langfuse_handler.langfuse_client.flush()
+    langfuse_handler.client.flush()
     sleep(2)
 
     trace = get_api().trace.get(trace_id=trace_id)
@@ -1020,7 +1020,7 @@ def test_link_langfuse_prompts_stream():
         for chunk in stream:
             output += chunk
 
-    langfuse_handler.langfuse_client.flush()
+    langfuse_handler.client.flush()
     sleep(2)
 
     trace = get_api().trace.get(trace_id=trace_id)
@@ -1106,7 +1106,7 @@ def test_link_langfuse_prompts_batch():
             },
         )
 
-    langfuse_handler.langfuse_client.flush()
+    langfuse_handler.client.flush()
 
     traces = get_api().trace.list(name=trace_name).data
 
@@ -1229,13 +1229,13 @@ class GetWeather(BaseModel):
         }
     ]
 
-    with handler.langfuse_client.start_as_current_span(
+    with handler.client.start_as_current_span(
         name="test_callback_openai_functions_with_tools"
     ) as span:
         trace_id = span.trace_id
         llm.bind_tools([address_tool, weather_tool]).invoke(messages)
 
-    handler.langfuse_client.flush()
+    handler.client.flush()
 
     trace = get_api().trace.get(trace_id=trace_id)
 
@@ -1328,11 +1328,11 @@ def test_multimodal():
         ],
     )
 
-    with handler.langfuse_client.start_as_current_span(name="test_multimodal") as span:
+    with handler.client.start_as_current_span(name="test_multimodal") as span:
         trace_id = span.trace_id
         model.invoke([message], config={"callbacks": [handler]})
 
-    handler.langfuse_client.flush()
+    handler.client.flush()
 
     trace = get_api().trace.get(trace_id=trace_id)
 
@@ -1417,14 +1417,14 @@ def call_model(state: MessagesState):
     handler = CallbackHandler()
 
     # Use the Runnable
-    with handler.langfuse_client.start_as_current_span(name="test_langgraph") as span:
+    with handler.client.start_as_current_span(name="test_langgraph") as span:
         trace_id = span.trace_id
         final_state = app.invoke(
             {"messages": [HumanMessage(content="what is the weather in sf")]},
             config={"configurable": {"thread_id": 42}, "callbacks": [handler]},
         )
     print(final_state["messages"][-1].content)
-    handler.langfuse_client.flush()
+    handler.client.flush()
 
     trace = get_api().trace.get(trace_id=trace_id)
 
@@ -1467,7 +1467,7 @@ def test_cached_token_usage():
     # invoke again to force cached token usage
     chain.invoke({"test_param": "in a funny way"}, config)
 
-    handler.langfuse_client.flush()
+    handler.client.flush()
 
     trace = get_api().trace.get(handler.get_trace_id())