From 352ea257e8e2804e614c5343503183789291eb3b Mon Sep 17 00:00:00 2001
From: Vasilije <8619304+Vasilije1990@users.noreply.github.com>
Date: Sat, 18 May 2024 19:11:05 +0200
Subject: [PATCH] test with a simple file

---
 .../infrastructure/llm/generic_llm_api/adapter.py  | 14 ++++++++++++++
 cognee/infrastructure/llm/openai/adapter.py        |  4 ++++
 cognee/shared/data_models.py                       |  1 +
 3 files changed, 19 insertions(+)

diff --git a/cognee/infrastructure/llm/generic_llm_api/adapter.py b/cognee/infrastructure/llm/generic_llm_api/adapter.py
index 2ee65178c..03a61a742 100644
--- a/cognee/infrastructure/llm/generic_llm_api/adapter.py
+++ b/cognee/infrastructure/llm/generic_llm_api/adapter.py
@@ -7,9 +7,23 @@ from tenacity import retry, stop_after_attempt
 from openai import AsyncOpenAI
 import openai
 
+from cognee.config import Config
 from cognee.infrastructure import infrastructure_config
 from cognee.infrastructure.llm.llm_interface import LLMInterface
 from cognee.infrastructure.llm.prompts import read_query_prompt
+from cognee.shared.data_models import MonitoringTool
+
+config = Config()
+config.load()
+
+if config.monitoring_tool == MonitoringTool.LANGFUSE:
+    from langfuse.openai import AsyncOpenAI, OpenAI
+elif config.monitoring_tool == MonitoringTool.LANGSMITH:
+    from langsmith import wrap_openai
+    from openai import AsyncOpenAI
+    AsyncOpenAI = wrap_openai(AsyncOpenAI())
+else:
+    from openai import AsyncOpenAI, OpenAI
 
 class GenericAPIAdapter(LLMInterface):
     """Adapter for Generic API LLM provider API """
diff --git a/cognee/infrastructure/llm/openai/adapter.py b/cognee/infrastructure/llm/openai/adapter.py
index 5c45b98bc..866f01357 100644
--- a/cognee/infrastructure/llm/openai/adapter.py
+++ b/cognee/infrastructure/llm/openai/adapter.py
@@ -15,6 +15,10 @@ config.load()
 
 if config.monitoring_tool == MonitoringTool.LANGFUSE:
     from langfuse.openai import AsyncOpenAI, OpenAI
+elif config.monitoring_tool == MonitoringTool.LANGSMITH:
+    from langsmith import wrap_openai
+    from openai import AsyncOpenAI
+    AsyncOpenAI = wrap_openai(AsyncOpenAI())
 else:
     from openai import AsyncOpenAI, OpenAI
 
diff --git a/cognee/shared/data_models.py b/cognee/shared/data_models.py
index 28fe72040..51b0124ba 100644
--- a/cognee/shared/data_models.py
+++ b/cognee/shared/data_models.py
@@ -257,4 +257,5 @@ class MonitoringTool(str, Enum):
     """ Monitoring tools """
     LANGFUSE = "langfuse"
     LLMLITE = "llmlite"
+    LANGSMITH = "langsmith"