add types to openai

Author: rohit-kadhe

src/langtrace_python_sdk/instrumentation/anthropic/types.py

@@ -1,3 +1,16 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 from typing import Dict, List, Optional, Iterator, TypedDict
 class MessagesCreateKwargs(TypedDict, total=False):
     system: str

src/langtrace_python_sdk/instrumentation/openai/instrumentation.py

@@ -1,25 +1,22 @@
 """
 Copyright (c) 2024 Scale3 Labs
-
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
-
     http://www.apache.org/licenses/LICENSE-2.0
-
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from typing import Collection, Optional, Any
 import importlib.metadata
 import logging
-from typing import Collection
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
-from opentelemetry.trace import get_tracer
+from opentelemetry.trace import get_tracer, TracerProvider
 from wrapt import wrap_function_wrapper
 
 from langtrace_python_sdk.instrumentation.openai.patch import (
@@ -35,59 +32,57 @@
 logging.basicConfig(level=logging.FATAL)
 
 
-class OpenAIInstrumentation(BaseInstrumentor):
+class OpenAIInstrumentation(BaseInstrumentor): # type: ignore
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return ["openai >= 0.27.0", "trace-attributes >= 4.0.5"]
 
-    def _instrument(self, **kwargs):
-        tracer_provider = kwargs.get("tracer_provider")
+    def _instrument(self, **kwargs: Any) -> None:
+        tracer_provider: Optional[TracerProvider] = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, "", tracer_provider)
-        version = importlib.metadata.version("openai")
+        version: str = importlib.metadata.version("openai")
 
         wrap_function_wrapper(
             "openai.resources.chat.completions",
             "Completions.create",
-            chat_completions_create("openai.chat.completions.create", version, tracer),
+            chat_completions_create(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.chat.completions",
             "AsyncCompletions.create",
-            async_chat_completions_create(
-                "openai.chat.completions.create_stream", version, tracer
-            ),
+            async_chat_completions_create(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.images",
             "Images.generate",
-            images_generate("openai.images.generate", version, tracer),
+            images_generate(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.images",
             "AsyncImages.generate",
-            async_images_generate("openai.images.generate", version, tracer),
+            async_images_generate(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.images",
             "Images.edit",
-            images_edit("openai.images.edit", version, tracer),
+            images_edit(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.embeddings",
             "Embeddings.create",
-            embeddings_create("openai.embeddings.create", version, tracer),
+            embeddings_create(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.embeddings",
             "AsyncEmbeddings.create",
-            async_embeddings_create("openai.embeddings.create", version, tracer),
+            async_embeddings_create(version, tracer),
         )
 
-    def _uninstrument(self, **kwargs):
+    def _uninstrument(self, **kwargs: Any) -> None:
         pass