fix(python-sdk): reconciliation in async context (#774)

packages/python-sdk/examples/langchain/v0.0.3/message-reconciliation.py

@@ -0,0 +1,81 @@
+from graph_rag_example_helpers.datasets.animals import fetch_documents
+animals = fetch_documents()
+
+from langchain_openai import OpenAIEmbeddings
+
+embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
+
+from langchain_core.vectorstores import InMemoryVectorStore
+
+vector_store = InMemoryVectorStore.from_documents(
+    documents=animals,
+    embedding=embeddings,
+)
+
+
+from graph_retriever.strategies import Eager
+from langchain_graph_retriever import GraphRetriever
+
+traversal_retriever = GraphRetriever(
+    store = vector_store,
+    edges = [("habitat", "habitat"), ("origin", "origin")],
+    strategy = Eager(k=5, start_k=1, max_depth=2),
+)
+
+
+results = traversal_retriever.invoke("what animals could be found near a capybara?")
+
+
+from langchain.chat_models import init_chat_model
+
+llm = init_chat_model("gpt-4o-mini", model_provider="openai")
+
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough, RunnableConfig
+import lunary
+
+handler = lunary.LunaryCallbackHandler()
+
+prompt = ChatPromptTemplate.from_template(
+"""Answer the question based only on the context provided.
+
+Context: {context}
+
+Question: {question}"""
+)
+
+def format_docs(docs):
+    return "\n\n".join(f"text: {doc.page_content} metadata: {doc.metadata}" for doc in docs)
+
+chain = (
+    {"context": traversal_retriever | format_docs, "question": RunnablePassthrough()}
+    | prompt
+    | llm
+    | StrOutputParser()
+)
+
+thread1 = lunary.open_thread()
+thread2 = lunary.open_thread()
+
+
+
+
+async def main():
+  msg_id = thread1.track_message({"role": "user", "content": "what animals could be found near a capybar"})
+  with lunary.parent(msg_id):
+    res = await chain.ainvoke("what animals could be found near a capybara?", RunnableConfig({"callbacks": [handler]}))
+    msg_id = thread1.track_message({"role": "assistant", "content": res})
+
+def main2():
+  msg_id = thread2.track_message({"role": "user", "content": "what animals could be found near a capybar"})
+  with lunary.parent(msg_id):
+    res = chain.invoke("what animals could be found near a capybara?", RunnableConfig({"callbacks": [handler]}))
+    msg_id = thread2.track_message({"role": "assistant", "content": res})
+
+
+import asyncio
+asyncio.run(main())
+main2()
+
+

packages/python-sdk/lunary/__init__.py

@@ -28,7 +28,7 @@
     identify,
 )  # DO NOT REMOVE `identify`` import
 from .tags import tags_ctx, tags  # DO NOT REMOVE `tags` import
-from .parent import parent_ctx, parent, get_parent  # DO NOT REMOVE `parent` import
+from .parent import parent_ctx, parent # DO NOT REMOVE `parent` import
 from .project import project_ctx  # DO NOT REMOVE `project` import
 
 logging.basicConfig()
@@ -49,14 +49,6 @@ class LunaryException(Exception):
 
 jsonpickle.handlers.register(BaseModel, PydanticHandler, base=True)
 
-def get_parent():
-    parent = parent_ctx.get()
-    if parent and parent.get("retrieved", False) == False:
-        parent_ctx.set({"message_id": parent["message_id"], "retrieved": True})
-        return parent.get("message_id", None)
-    return None
-
-
 def config(
     app_id: str | None = None,
     verbose: str | None = None,
@@ -70,8 +62,8 @@ def get_parent_run_id(parent_run_id: str, run_type: str, app_id: str, run_id: st
     if parent_run_id == "None":
         parent_run_id = None
 
-    parent_from_ctx = get_parent()
-    if parent_from_ctx and run_type != "thread":
+    parent_from_ctx = parent_ctx.get().get("message_id") if parent_ctx.get() else None
+    if not parent_run_id and parent_from_ctx and run_type != "thread":
         return str(create_uuid_from_string(str(parent_from_ctx) + str(app_id)))
 
     if parent_run_id:

packages/python-sdk/lunary/parent.py

@@ -16,10 +16,3 @@ def __exit__(self, exc_type, exc_value, exc_tb):
 
 def parent(id: str) -> ParentContextManager:
     return ParentContextManager(id)
-
-def get_parent():
-  parent = parent_ctx.get()
-  if parent and parent.get("retrieved", False) == False:
-    parent_ctx.set({"message_id": parent["message_id"], "retrieved": True})
-    return parent.get("message_id", None)
-  return None

packages/python-sdk/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lunary"
-version = "1.4.5"
+version = "1.4.6"
 description = "Observability, analytics and evaluations for AI agents and chatbots."
 authors = ["lunary <[email protected]>"]
 readme = "README.md"