From 5ac7d3d588d247bc6f5a399a74dd78291ea24375 Mon Sep 17 00:00:00 2001 From: csgulati09 Date: Tue, 13 Aug 2024 02:27:01 +0530 Subject: [PATCH 1/4] fix: handling None type payload --- .../llama_index/portkey_llama_callback.py | 42 +++++++++++-------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/portkey_ai/llms/llama_index/portkey_llama_callback.py b/portkey_ai/llms/llama_index/portkey_llama_callback.py index 0d5c77a..569fcbb 100644 --- a/portkey_ai/llms/llama_index/portkey_llama_callback.py +++ b/portkey_ai/llms/llama_index/portkey_llama_callback.py @@ -126,22 +126,31 @@ def on_event_end( """Run when an event ends.""" span_id = event_id - if event_type == "llm": - response_payload = self.llm_event_end(payload, event_id) - elif event_type == "embedding": - response_payload = self.embedding_event_end(payload, event_id) - elif event_type == "agent_step": - response_payload = self.agent_step_event_end(payload, event_id) - elif event_type == "function_call": - response_payload = self.function_call_event_end(payload, event_id) - elif event_type == "query": - response_payload = self.query_event_end(payload, event_id) - elif event_type == "retrieve": - response_payload = self.retrieve_event_end(payload, event_id) - elif event_type == "templating": - response_payload = self.templating_event_end(payload, event_id) + if payload is None: + response_payload = {} + if span_id in self.event_map: + event = self.event_map[event_id] + start_time = event["start_time"] + end_time = int(datetime.now().timestamp()) + total_time = (end_time - start_time) * 1000 + response_payload["response_time"] = total_time else: - response_payload = payload + if event_type == "llm": + response_payload = self.llm_event_end(payload, event_id) + elif event_type == "embedding": + response_payload = self.embedding_event_end(payload, event_id) + elif event_type == "agent_step": + response_payload = self.agent_step_event_end(payload, event_id) + elif event_type == "function_call": + response_payload = self.function_call_event_end(payload, event_id) + elif event_type == "query": + response_payload = self.query_event_end(payload, event_id) + elif event_type == "retrieve": + response_payload = self.retrieve_event_end(payload, event_id) + elif event_type == "templating": + response_payload = self.templating_event_end(payload, event_id) + else: + response_payload = payload self.event_map[span_id]["response"] = response_payload @@ -149,7 +158,6 @@ def on_event_end( def start_trace(self, trace_id: Optional[str] = None) -> None: """Run when an overall trace is launched.""" - if trace_id == "index_construction": self.global_trace_id = self.metadata.get("traceId", str(uuid4())) # type: ignore [union-attr] @@ -230,7 +238,7 @@ def llm_event_end(self, payload: Any, event_id) -> Any: ) self.response["body"].update({"id": event_id}) self.response["body"].update({"created": int(time.time())}) - self.response["body"].update({"model": data.raw.get("model", "")}) + self.response["body"].update({"model": getattr(data, "model", "")}) self.response["headers"] = {} self.response["streamingMode"] = self.streamingMode From 92dc64c24d9e1b1fcf949953690e659e4f219b67 Mon Sep 17 00:00:00 2001 From: csgulati09 Date: Tue, 13 Aug 2024 14:30:52 +0530 Subject: [PATCH 2/4] fix: import style changed --- portkey_ai/__init__.py | 7 +++++++ portkey_ai/llms/langchain/__init__.py | 4 ++-- portkey_ai/llms/langchain/portkey_langchain_callback.py | 2 +- portkey_ai/llms/llama_index/__init__.py | 4 ++-- portkey_ai/llms/llama_index/portkey_llama_callback.py | 2 +- tests/test_llm_langchain.py | 4 ++-- tests/test_llm_llamaindex.py | 4 ++-- 7 files changed, 17 insertions(+), 10 deletions(-) diff --git a/portkey_ai/__init__.py b/portkey_ai/__init__.py index 2f325ab..c093ee7 100644 --- a/portkey_ai/__init__.py +++ b/portkey_ai/__init__.py @@ -73,6 +73,11 @@ PORTKEY_GATEWAY_URL, ) +from portkey_ai.llms.langchain.portkey_langchain_callback import ( + LangchainCallbackHandler, +) +from portkey_ai.llms.llama_index.portkey_llama_callback import LlamaIndexCallbackHandler + api_key = os.environ.get(PORTKEY_API_KEY_ENV) base_url = os.environ.get(PORTKEY_PROXY_ENV, PORTKEY_BASE_URL) config: Optional[Union[Mapping, str]] = None @@ -144,5 +149,7 @@ "AsyncVectorStores", "AsyncVectorFiles", "AsyncVectorFileBatches", + "LangchainCallbackHandler", + "LlamaIndexCallbackHandler", "openai", ] diff --git a/portkey_ai/llms/langchain/__init__.py b/portkey_ai/llms/langchain/__init__.py index de0f7f0..b6d71b8 100644 --- a/portkey_ai/llms/langchain/__init__.py +++ b/portkey_ai/llms/langchain/__init__.py @@ -1,5 +1,5 @@ from .chat import ChatPortkey from .completion import PortkeyLLM -from .portkey_langchain_callback import PortkeyLangchain +from .portkey_langchain_callback import LangchainCallbackHandler -__all__ = ["ChatPortkey", "PortkeyLLM", "PortkeyLangchain"] +__all__ = ["ChatPortkey", "PortkeyLLM", "LangchainCallbackHandler"] diff --git a/portkey_ai/llms/langchain/portkey_langchain_callback.py b/portkey_ai/llms/langchain/portkey_langchain_callback.py index c45864f..7d2899d 100644 --- a/portkey_ai/llms/langchain/portkey_langchain_callback.py +++ b/portkey_ai/llms/langchain/portkey_langchain_callback.py @@ -13,7 +13,7 @@ raise ImportError("Please pip install langchain-core to use PortkeyLangchain") -class PortkeyLangchain(BaseCallbackHandler): +class LangchainCallbackHandler(BaseCallbackHandler): def __init__( self, api_key: str, diff --git a/portkey_ai/llms/llama_index/__init__.py b/portkey_ai/llms/llama_index/__init__.py index 9530d3e..f9f92d0 100644 --- a/portkey_ai/llms/llama_index/__init__.py +++ b/portkey_ai/llms/llama_index/__init__.py @@ -1,3 +1,3 @@ -from .portkey_llama_callback import PortkeyLlamaindex +from .portkey_llama_callback import LlamaIndexCallbackHandler -__all__ = ["PortkeyLlamaindex"] +__all__ = ["LlamaIndexCallbackHandler"] diff --git a/portkey_ai/llms/llama_index/portkey_llama_callback.py b/portkey_ai/llms/llama_index/portkey_llama_callback.py index 569fcbb..615fcfe 100644 --- a/portkey_ai/llms/llama_index/portkey_llama_callback.py +++ b/portkey_ai/llms/llama_index/portkey_llama_callback.py @@ -23,7 +23,7 @@ raise ImportError("Please pip install llama-index to use Portkey Callback Handler") -class PortkeyLlamaindex(LlamaIndexBaseCallbackHandler): +class LlamaIndexCallbackHandler(LlamaIndexBaseCallbackHandler): def __init__( self, api_key: str, diff --git a/tests/test_llm_langchain.py b/tests/test_llm_langchain.py index 3441d85..a15e129 100644 --- a/tests/test_llm_langchain.py +++ b/tests/test_llm_langchain.py @@ -6,7 +6,7 @@ import pytest from tests.utils import read_json_file -from portkey_ai.llms.langchain import PortkeyLangchain +from portkey_ai import LangchainCallbackHandler from langchain.chat_models import ChatOpenAI from langchain_core.prompts import ChatPromptTemplate from langchain.chains import LLMChain @@ -15,7 +15,7 @@ class TestLLMLangchain: - client = PortkeyLangchain + client = LangchainCallbackHandler parametrize = pytest.mark.parametrize("client", [client], ids=["strict"]) models = read_json_file("./tests/models.json") diff --git a/tests/test_llm_llamaindex.py b/tests/test_llm_llamaindex.py index 959dd81..f4a51e1 100644 --- a/tests/test_llm_llamaindex.py +++ b/tests/test_llm_llamaindex.py @@ -6,7 +6,7 @@ import pytest from tests.utils import read_json_file -from portkey_ai.llms.llama_index import PortkeyLlamaindex +from portkey_ai import LlamaIndexCallbackHandler from llama_index.llms.openai import OpenAI @@ -24,7 +24,7 @@ class TestLLMLlamaindex: - client = PortkeyLlamaindex + client = LlamaIndexCallbackHandler parametrize = pytest.mark.parametrize("client", [client], ids=["strict"]) models = read_json_file("./tests/models.json") From d3c250c356e64d1c3dd6eab176ff92cd40a01c7e Mon Sep 17 00:00:00 2001 From: csgulati09 Date: Tue, 13 Aug 2024 15:13:51 +0530 Subject: [PATCH 3/4] feat: new folder structure and imports --- portkey_ai/__init__.py | 9 ++------- portkey_ai/langchain/__init__.py | 3 +++ .../portkey_langchain_callback_handler.py} | 0 portkey_ai/llamaindex/__init__.py | 3 +++ .../portkey_llama_callback_handler.py} | 0 portkey_ai/llms/langchain/__init__.py | 6 ++++-- portkey_ai/llms/llama_index/__init__.py | 3 --- tests/test_llm_langchain.py | 2 +- tests/test_llm_llamaindex.py | 2 +- 9 files changed, 14 insertions(+), 14 deletions(-) create mode 100644 portkey_ai/langchain/__init__.py rename portkey_ai/{llms/langchain/portkey_langchain_callback.py => langchain/portkey_langchain_callback_handler.py} (100%) create mode 100644 portkey_ai/llamaindex/__init__.py rename portkey_ai/{llms/llama_index/portkey_llama_callback.py => llamaindex/portkey_llama_callback_handler.py} (100%) diff --git a/portkey_ai/__init__.py b/portkey_ai/__init__.py index c093ee7..9271255 100644 --- a/portkey_ai/__init__.py +++ b/portkey_ai/__init__.py @@ -73,11 +73,6 @@ PORTKEY_GATEWAY_URL, ) -from portkey_ai.llms.langchain.portkey_langchain_callback import ( - LangchainCallbackHandler, -) -from portkey_ai.llms.llama_index.portkey_llama_callback import LlamaIndexCallbackHandler - api_key = os.environ.get(PORTKEY_API_KEY_ENV) base_url = os.environ.get(PORTKEY_PROXY_ENV, PORTKEY_BASE_URL) config: Optional[Union[Mapping, str]] = None @@ -149,7 +144,7 @@ "AsyncVectorStores", "AsyncVectorFiles", "AsyncVectorFileBatches", - "LangchainCallbackHandler", - "LlamaIndexCallbackHandler", + # "LangchainCallbackHandler", + # "LlamaIndexCallbackHandler", "openai", ] diff --git a/portkey_ai/langchain/__init__.py b/portkey_ai/langchain/__init__.py new file mode 100644 index 0000000..4ee8054 --- /dev/null +++ b/portkey_ai/langchain/__init__.py @@ -0,0 +1,3 @@ +from .portkey_langchain_callback_handler import LangchainCallbackHandler + +__all__ = ["LangchainCallbackHandler"] diff --git a/portkey_ai/llms/langchain/portkey_langchain_callback.py b/portkey_ai/langchain/portkey_langchain_callback_handler.py similarity index 100% rename from portkey_ai/llms/langchain/portkey_langchain_callback.py rename to portkey_ai/langchain/portkey_langchain_callback_handler.py diff --git a/portkey_ai/llamaindex/__init__.py b/portkey_ai/llamaindex/__init__.py new file mode 100644 index 0000000..55e48a8 --- /dev/null +++ b/portkey_ai/llamaindex/__init__.py @@ -0,0 +1,3 @@ +from .portkey_llama_callback_handler import LlamaIndexCallbackHandler + +__all__ = ["LlamaIndexCallbackHandler"] diff --git a/portkey_ai/llms/llama_index/portkey_llama_callback.py b/portkey_ai/llamaindex/portkey_llama_callback_handler.py similarity index 100% rename from portkey_ai/llms/llama_index/portkey_llama_callback.py rename to portkey_ai/llamaindex/portkey_llama_callback_handler.py diff --git a/portkey_ai/llms/langchain/__init__.py b/portkey_ai/llms/langchain/__init__.py index b6d71b8..07e0ff1 100644 --- a/portkey_ai/llms/langchain/__init__.py +++ b/portkey_ai/llms/langchain/__init__.py @@ -1,5 +1,7 @@ from .chat import ChatPortkey from .completion import PortkeyLLM -from .portkey_langchain_callback import LangchainCallbackHandler -__all__ = ["ChatPortkey", "PortkeyLLM", "LangchainCallbackHandler"] +__all__ = [ + "ChatPortkey", + "PortkeyLLM", +] diff --git a/portkey_ai/llms/llama_index/__init__.py b/portkey_ai/llms/llama_index/__init__.py index f9f92d0..e69de29 100644 --- a/portkey_ai/llms/llama_index/__init__.py +++ b/portkey_ai/llms/llama_index/__init__.py @@ -1,3 +0,0 @@ -from .portkey_llama_callback import LlamaIndexCallbackHandler - -__all__ = ["LlamaIndexCallbackHandler"] diff --git a/tests/test_llm_langchain.py b/tests/test_llm_langchain.py index a15e129..8cfe75e 100644 --- a/tests/test_llm_langchain.py +++ b/tests/test_llm_langchain.py @@ -6,7 +6,7 @@ import pytest from tests.utils import read_json_file -from portkey_ai import LangchainCallbackHandler +from portkey_ai.langchain import LangchainCallbackHandler from langchain.chat_models import ChatOpenAI from langchain_core.prompts import ChatPromptTemplate from langchain.chains import LLMChain diff --git a/tests/test_llm_llamaindex.py b/tests/test_llm_llamaindex.py index f4a51e1..a0e9fff 100644 --- a/tests/test_llm_llamaindex.py +++ b/tests/test_llm_llamaindex.py @@ -6,7 +6,7 @@ import pytest from tests.utils import read_json_file -from portkey_ai import LlamaIndexCallbackHandler +from portkey_ai.llamaindex import LlamaIndexCallbackHandler from llama_index.llms.openai import OpenAI From 5165a6f1b56cb5657ea955d1e867ec8ae4da997d Mon Sep 17 00:00:00 2001 From: csgulati09 Date: Tue, 13 Aug 2024 15:19:22 +0530 Subject: [PATCH 4/4] fix: remove comments --- portkey_ai/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/portkey_ai/__init__.py b/portkey_ai/__init__.py index 9271255..2f325ab 100644 --- a/portkey_ai/__init__.py +++ b/portkey_ai/__init__.py @@ -144,7 +144,5 @@ "AsyncVectorStores", "AsyncVectorFiles", "AsyncVectorFileBatches", - # "LangchainCallbackHandler", - # "LlamaIndexCallbackHandler", "openai", ]