Fix LangChain session id generation

Author: dmelikyan

.github/workflows/daily-tests.yml

@@ -32,5 +32,4 @@ jobs:
         run: poetry update
 
       - name: Run unit tests
-        continue-on-error: true
         run: poetry run test

graphsignal/callbacks/langchain/v2.py

@@ -27,8 +27,10 @@ def __init__(self, **kwargs: Any) -> None:
         self._span_map = {}
 
         # create session_id, if not passed
+        self._generated_session_id = False
         if 'session_id' not in self._passed_tags and 'session_id' not in self._context_tags:
             self._passed_tags['session_id'] = uuid_sha1(size=12)
+            self._generated_session_id = True
 
     def _start_trace(self, parent_run_id, run_id, operation):
         # propagate context tags just in case contextvars are not propagated
@@ -58,7 +60,7 @@ def _stop_trace(self, run_id):
             span.stop()
 
     def on_llm_start(
-            self, 
+            self,
             serialized: Dict[str, Any], 
             prompts: List[str], 
             *,
@@ -618,7 +620,8 @@ async def on_agent_finish(
             self,
             finish: AgentFinish,
             **kwargs: Any) -> None:
-        pass
+        if self._generated_session_id:
+            graphsignal.set_context_tag('session_id', None)
 
     async def on_text(
             self,

graphsignal/version.py

@@ -1 +1 @@
-__version__ = '0.15.5'
+__version__ = '0.15.6'

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphsignal"
-version = "0.15.5"
+version = "0.15.6"
 description = "Graphsignal Tracer for Python"
 authors = ["Graphsignal, Inc. <devops@graphsignal.com>"]
 license = "Apache-2.0"
@@ -37,8 +37,9 @@ langchain = "*"
 langchainhub = "*"
 numexpr = "*"
 llama_index = "*"
-llama-index-llms-langchain = "*"
 nltk = "*"
+langchain-community = "*"
+llama-index-llms-langchain = "*"
 
 [tool.poetry.scripts]
 update-client = "tools.update_client:main"

pyproject.toml

@@ -68,6 +68,8 @@ async def test_callback_tags(self, mocked_upload_span):
     @patch.object(Uploader, 'upload_span')
     @patch('graphsignal.callbacks.langchain.v2.uuid_sha1', return_value='s1')
     async def test_chain(self, mocked_uuid_sha1, mocked_upload_span):
+        os.environ["LANGCHAIN_TRACING_V2"] = "false"
+
         graphsignal.set_context_tag('ct1', 'v1')
 
         llm = DummyLLM()
@@ -96,6 +98,8 @@ def find_span(op_name):
     @patch.object(Uploader, 'upload_span')
     @patch('graphsignal.callbacks.langchain.v2.uuid_sha1', return_value='s1')
     async def test_chain_async(self, mocked_uuid_sha1, mocked_upload_span):
+        os.environ["LANGCHAIN_TRACING_V2"] = "false"
+
         graphsignal.set_context_tag('ct1', 'v1')
 
         llm = DummyLLM()

test/callbacks/langchain/test_v2.py

@@ -28,6 +28,7 @@
 from openai.types.create_embedding_response import Usage
 
 os.environ['OPENAI_API_KEY'] = 'sk-kkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkk'
+os.environ['LANGCHAIN_API_KEY'] = 'kkk'
 
 class OpenAIRecorderTest(unittest.IsolatedAsyncioTestCase):
     async def asyncSetUp(self):
@@ -410,6 +411,7 @@ async def test_ret_gen():
             self.assertEqual(output_json, [{
                 'choices': [{'delta': {'content': '',
                                         'function_call': None,
+                                        'refusal': None,
                                         'role': 'assistant',
                                         'tool_calls': None},
                             'finish_reason': 'stop',
@@ -419,9 +421,12 @@ async def test_ret_gen():
                 'id': 'chatcmpl-8OQctcUZdbOT1KO8x2IKGUQE7G33b',
                 'model': 'gpt-4-0613',
                 'object': 'chat.completion.chunk',
-                'system_fingerprint': None},
+                'service_tier': None,
+                'system_fingerprint': None,
+                'usage': None},
                 {'choices': [{'delta': {'content': 'We',
                                         'function_call': None,
+                                        'refusal': None,
                                         'role': None,
                                         'tool_calls': None},
                             'finish_reason': 'stop',
@@ -431,7 +436,9 @@ async def test_ret_gen():
                 'id': 'chatcmpl-8OQctcUZdbOT1KO8x2IKGUQE7G33b',
                 'model': 'gpt-4-0613',
                 'object': 'chat.completion.chunk',
-                'system_fingerprint': None}])
+                'service_tier': None,
+                'system_fingerprint': None,
+                'usage': None}])
 
     @patch.object(Uploader, 'upload_span')
     async def test_embedding_create(self, mocked_upload_span):