cavisson-pythonagent 0.0.1__py3-none-any.whl

pythonagent/agent/probes/langchain/langchain_community.py

@@ -0,0 +1,136 @@
+
+from ..base import ExitCallInterceptor
+from ..span import RetrievalSpan
+from pythonagent.utils import  get_current_timestamp_in_us, get_validated_duration, split_fp_topo
+import logging
+#from .constants import text_embedding_models, vectorstore_classes
+MAX_LENGTH = 100
+
+"""
+A wrapper class to log and monitor retrieval calls 
+made to vector databases.It wraps Langchain Community 
+retrievers to intercept _similarity_search calls.
+Tracks input queries and number of documents returned.
+"""
+
+class LangChainCommunityVectorStoreInterceptor(ExitCallInterceptor):
+
+    def __init__(self, agent, cls, vectordb_type):
+        super().__init__(agent, cls)
+        self.vectordb_type = vectordb_type
+
+    def _similarity_search(self, similarity_search, *args, **kwargs):
+
+        start_time, exit_call, bt, fqm, host_name_str, context = 0, 0, None, None, None, None
+        test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = None, None, None, None, None, None, None
+
+
+        try:
+            bt, callout_type, db = self.bt, "VECTORDB", self.vectordb_type
+            host, port, query, username, query_params = "NA", "NA", "NA", "NA", None
+            fqm = f"{similarity_search.__module__}.{similarity_search.__qualname__}"
+            host_name_str = f"NA|{host}|{port}|NA|{callout_type}|{db}|NA|NA|NA|{username}|NA"
+            context = self.agent.get_transaction_context()
+            if context:
+                context.host_str = host_name_str
+                context.entry_point_fqm = fqm
+
+            start_time = get_current_timestamp_in_us()
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info("langchain community vectorstore db call begin bt {} host {} query {} query_params {} start_time {}".format(bt, host, query, query_params, start_time))
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time)
+
+        except Exception as e:
+            self.agent.logger.exception("Exception in langchain community vectorstore db call begin {}".format(e))
+
+        try:
+            result = similarity_search(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("Error in application: {}".format(e))
+            try:
+                end_time = get_current_timestamp_in_us()
+                duration = get_validated_duration(start_time, end_time, "langchain_community_vectorstore")
+                if context:
+                    context.status_code = 500
+                    context.db_callout_time = 0
+                if self.agent.logger.isEnabledFor(logging.INFO):
+                    self.agent.logger.info(
+                        "langchain_community_vectorstore db call end bt {} exit_call {} duration {} start_time {} fqm {}host_name_str {}".format(
+                            bt, exit_call, duration, start_time, fqm, host_name_str))
+                self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+            except Exception as e:
+                self.agent.logger.exception("Exception in Original Method DB Call End {}".format(e))
+
+            raise
+
+        try:
+            end_time = get_current_timestamp_in_us()
+            duration = get_validated_duration(start_time, end_time, "langchain_community_vectorstore")
+            if context:
+                context.status_code = 200
+                context.db_callout_time = 0
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info("langchain community vectordb db call end bt {} exit_call {} duration {} start_time {} fqm {}host_name_str {}".format(bt, exit_call, duration, start_time, fqm, host_name_str))
+            self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+
+
+            if self.agent.nf_enabled:
+                # Capturing input prompt and output response, formatting it
+                output_docs = []
+                for document in result:
+                    doc_info = {
+                        "source": document.metadata["source"],
+                        "content": document.page_content[:MAX_LENGTH]
+                    }
+                    output_docs.append(doc_info)
+
+
+                vectorstore_object = args[0]
+                query = args[1]
+
+                meta = {
+                    "span_kind": "retrieval",
+                    "input": str(query),
+                    "output": {
+                        "num_of_docs": len(result),
+                        "doc_sample": output_docs
+                    }
+                }
+
+                if context:
+                    fp_topo = context.fp_topo
+                    test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = split_fp_topo(fp_topo)
+
+                retrieval_span = RetrievalSpan(fqm, start_time, duration, "ok", meta, test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp)
+                retrieval_span_json = retrieval_span.create_span_json()
+
+                # SDK API call to send logs to NF
+                if self.agent.logger.isEnabledFor(logging.INFO):
+                     self.agent.logger.info("langchain community vectordb sending vectordb span {}".format(retrieval_span_json))
+                self.agent.dump_log_message(bt, retrieval_span_json)
+
+        except Exception as e:
+            self.agent.logger.exception("Exception in langchain community vectorstore db call end {}".format(e))
+
+        return result
+
+class LangChainCommunityEmbeddingsInterceptor(ExitCallInterceptor):
+    def _embed_query(self, embed_query, *args, **kwargs):
+        result = embed_query(*args, **kwargs)
+        return result
+
+
+def intercept_langchain_community_vectorstores(agent, mod):
+    agent.logger.warning("Instrument module: langchain_community.vectorstores{}".format(", mod: {}".format(mod) if mod else ""))
+
+    LangChainCommunityVectorStoreInterceptor(agent, mod.Chroma, "Chroma").attach("similarity_search")
+
+    #for vectorstore_class in vectorstore_classes:
+    #    LangChainCommunityVectorStoreInterceptor(agent, mod.vectorstore_class).attach("similarity_search")
+
+def intercept_langchain_community_embeddings(agent, mod):
+    agent.logger.warning("Instrument module: langchain_community.embeddings{}".format(", mod: {}".format(mod) if mod else ""))
+    pass
+
+    #for text_embedding_model in text_embedding_models:
+    #    LangChainCommunityVectorStoreInterceptor(agent, mod.text_embedding_model).attach("embed_query")

pythonagent/agent/probes/langchain/langchain_core.py

@@ -0,0 +1,32 @@
+from logging import exception
+
+from click import prompt
+
+from ..base import ExitCallInterceptor
+from pythonagent.utils import  get_current_timestamp_in_us, get_validated_duration
+
+#from .base_llm import LangChainBaseLLMInterceptor
+#from .base_chat_model import LangChainBaseChatModelInterceptor
+#from .langchain_runnable_sequence import LangChainRunnableSequenceInterceptor
+from .base_tool import LangChainBaseToolInterceptor
+
+"""
+# Interceptor API for instrument langchain_core.language_models module
+def intercept_langchain_core_language_models(agent, mod):
+    langchain_base_llm_interceptor = LangChainBaseLLMInterceptor(agent, mod.llms.BaseLLM)
+    langchain_base_llm_interceptor.attach(["invoke", "generate", "stream"])
+
+    langchain_base_chat_model_interceptor = LangChainBaseChatModelInterceptor(agent, mod.chat_models.BaseChatModel)
+    langchain_base_chat_model_interceptor.attach(["invoke", "generate", "stream"])
+
+# Interceptor API for instrument langchain_core.runnables module
+def intercept_langchain_core_runnables(agent, mod):
+    langchain_runnable_sequence_interceptor = LangChainRunnableSequenceInterceptor(agent, mod.base.RunnableSequence)
+    langchain_runnable_sequence_interceptor.attach(["invoke", "batch", "stream"])
+"""
+
+# Interceptor API for instrument langchain_core.tools module
+def intercept_langchain_core_tools(agent, mod):
+    agent.logger.warning("Instrument module: langchain_core.tools{}".format(", mod: {}".format(mod) if mod else ""))
+    langchain_base_tool_interceptor = LangChainBaseToolInterceptor(agent, mod.base.BaseTool)
+    langchain_base_tool_interceptor.attach("invoke")

pythonagent/agent/probes/langchain/langchain_openai.py

@@ -0,0 +1,110 @@
+from ..base import ExitCallInterceptor
+from ..span import EmbeddingSpan
+from pythonagent.utils import get_current_timestamp_in_us, get_validated_duration, split_fp_topo
+import logging
+SAMPLE_LEN = 3
+
+"""
+A wrapper class to monitor calls to OpenAI Embeddings in LangChain.
+Intercepts _embed_query requests for logging or metrics.
+Tracks input, outputs, embedding model name and provider
+vector length, vector sample,.
+"""
+
+
+class LangChainOpenAIEmbeddingsInterceptor(ExitCallInterceptor):
+    def _embed_query(self, embed_query, *args, **kwargs):
+
+        start_time, exit_call, bt, embeddings_model, context, fqm, host_name_str = 0, 0, None, None, None, None, None
+        test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = None, None, None, None, None, None, None
+        try:
+            bt, callout_type, provider = self.bt, "EMBEDDINGS", "OpenAI"
+            host, port, query, username, query_params = "NA", "NA", "NA", "NA", None
+
+            openai_embeddings_object = args[0]
+            embeddings_model = openai_embeddings_object.model
+
+            db = provider + "_" + embeddings_model
+
+            fqm = f"{embed_query.__module__}.{embed_query.__qualname__}"
+            host_name_str = f"NA|{host}|{port}|NA|{callout_type}|{db}|NA|NA|NA|{username}|NA"
+
+            context = self.agent.get_transaction_context()
+            if context:
+                context.host_str = host_name_str
+                context.entry_point_fqm = fqm
+
+            start_time = get_current_timestamp_in_us()
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info("langchain openai embeddings db call begin bt {} host {} query {} query_params {} start_time {}".format(bt, host, query, query_params, start_time))
+
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time)
+
+        except Exception as e:
+            self.agent.logger.exception("langchain openai embeddings db call begin {}".format(e))
+
+        try:
+            result = embed_query(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("Error in application: {}".format(e))
+            try:
+                end_time = get_current_timestamp_in_us()
+                duration = get_validated_duration(start_time, end_time, "LangchainOpenAI Embeddings")
+                if context:
+                    context.status_code = 500
+                    context.db_callout_time = 0
+                if self.agent.logger.isEnabledFor(logging.INFO):
+                    self.agent.logger.info("langchain openai embeddings db call end bt {} exit_call {} duration {} start_time {} fqm {}host_name_str {}".format(bt, exit_call, duration, start_time, fqm, host_name_str))
+                self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+            except Exception as e:
+                self.agent.logger.exception("Exception in Original Method langchain openai embeddings {}".format(e))
+
+            raise
+
+        try:
+            end_time = get_current_timestamp_in_us()
+            duration = get_validated_duration(start_time, end_time, "LangchainOpenAI Embeddings")
+            if context:
+                context.status_code = 200
+                context.db_callout_time = 0
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info(
+                    "langchain openai embeddings db call end bt {} exit_call {} duration {} start_time {} fqm {}host_name_str {}".format(
+                        bt, exit_call, duration, start_time, fqm, host_name_str))
+            self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+
+            if self.agent.nf_enabled:
+                # Capturing input prompt and output response, formatting it
+                meta = {"span_kind": "embedding",
+                        "input": args[1],
+                        "output": {"vector_length": len(result),
+                                   "vector_sample": result[:SAMPLE_LEN],
+                                    "model_name": embeddings_model,
+                                    "model_provider": "OpenAI"
+                                   }
+                        }
+
+                if context:
+                    fp_topo = context.fp_topo
+                    test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = split_fp_topo(fp_topo)
+
+                embedding_span = EmbeddingSpan(fqm, start_time, duration, "ok", meta, test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp)
+                embedding_span_json = embedding_span.create_span_json()
+
+                # SDK API call to send logs to NF
+                if self.agent.logger.isEnabledFor(logging.INFO):
+                    self.agent.logger.info("langchain openai embeddings sending embeddings span {}".format(embedding_span_json))
+                self.agent.dump_log_message(bt, embedding_span_json)
+
+        except Exception as e:
+            self.agent.logger.exception("langchain openai embeddings db call end {}".format(e))
+
+
+
+        return result
+
+
+# Interceptor API to instrument langchain_openai module
+def intercept_langchain_openai(agent, mod):
+    agent.logger.warning("Instrument module: langchain_openai{}".format(", mod: {}".format(mod) if mod else ""))
+    LangChainOpenAIEmbeddingsInterceptor(agent, mod.embeddings.OpenAIEmbeddings).attach("embed_query")

pythonagent/agent/probes/logging/__init__.py

@@ -0,0 +1,106 @@
+
+"""Intercept the Python logging module.
+
+"""
+
+import logging
+import sys
+import time
+import traceback
+
+from ..base import BaseInterceptor
+
+
+class LoggingInterceptor(BaseInterceptor):
+
+    def _callHandlers(self, callHandlers, logger, record):
+
+        if record.name.startswith('pythonagent.agent'):
+            return callHandlers(logger, record)
+
+        #self.agent.application_loggers.add(logger)
+
+        try:
+            self.agent.logger.info("error message as {}".format(record.getMessage()))
+            self.add_logged_error(record)
+
+            if type(record.msg) is not str:
+                record.msg = str(record.msg)
+
+            context = self.agent.get_transaction_context()
+            if context:
+                if context.aws_request_id:
+                    record.msg = f"[AWS Lambda Request ID: {context.aws_request_id}] " + record.msg
+                if context.fp_topo:
+                    record.msg = record.msg + str(context.fp_topo)
+                self.agent.logger.debug("Modified log message {}".format(record.msg))
+
+            else:
+                self.agent.logger.info("Transaction context not found while appending application log in logging.__init__")
+
+        except Exception as e:
+            self.agent.logger.exception("Error in logging.__init__ _callHandlers {}".format(e))
+
+        callHandlers(logger, record)
+
+
+    def add_logged_error(self, record):
+        """Add an error to the BT if the log record should trigger an error.
+
+        If the level is greater or equal than the threshold and the message is
+        not ignored, create an ErrorInfo object and add it to the BT.
+
+        """
+        if record.levelno < 40:
+            return
+
+        try:
+            start_time = round(time.time() * 1000)
+            exception_class = ''
+            exception_cause = ''
+            exception_message = ''
+            file_name = ''
+            line_number = ''
+            throwing_method = ''
+            stack_trace_str = ''
+
+            if record.exc_info:
+                exception_class_type, exception_message, traceback_object = record.exc_info
+
+                exception_class = str(exception_class_type.__name__)
+                exception_cause = str(exception_class_type.__cause__)
+                exception_message = str(exception_message)
+
+                tb = traceback.extract_tb(traceback_object)[-1]  # Get last traceback frame
+
+                file_name, line_number, throwing_method, _ = tb
+
+                full_stack_trace_str = ''.join(traceback.format_exception(*record.exc_info))
+                full_stack_trace_str = full_stack_trace_str.strip()
+
+                stack_trace_list = full_stack_trace_str.split("\n")
+                size = 2 if sys.version_info < (3, 11, 0) else 3
+                max_depth = self.agent.exception_stack_depth * size
+
+                first_line = stack_trace_list[0]
+                last_line = stack_trace_list[-1]
+                stack_trace_list = stack_trace_list[1:-1]
+
+                stack_trace_len = len(stack_trace_list)
+                if max_depth < stack_trace_len:
+                    stack_trace_list = stack_trace_list[-max_depth:]
+
+                stack_trace_str = first_line + "\n" + "\n".join(stack_trace_list) + "\n" + last_line
+
+            self.agent.logger.debug('The Exception is generated with  Exception_class as {} ,Exception_message as {} ,Exception_cause as {}  and Throwing method as {} with a Depth of {} in file {}'.format(exception_class,exception_message,exception_cause,throwing_method,self.agent.exception_stack_depth,file_name))
+            self.agent.exceptiondump(self.bt, None, start_time, exception_class, exception_message, "None", throwing_method, exception_cause, line_number, stack_trace_str,file_name)
+
+        except Exception as e:
+            self.agent.logger.exception("Error in creating exception record: {}".format(e))
+
+        self.agent.logger.info('Logging application error as message {}'.format(record.msg))
+
+
+def intercept_logging(agent, mod):
+    agent.logger.warning("Instrument module: logging{}".format(", mod: {}".format(mod) if mod else ""))
+    LoggingInterceptor(agent, mod.Logger).attach('callHandlers')

pythonagent/agent/probes/message_brokers/__init__.py

@@ -0,0 +1,4 @@
+from .pika import intercept_pika
+
+__all__ = ['intercept_pika']
+

pythonagent/agent/probes/message_brokers/pika.py

@@ -0,0 +1,126 @@
+from ..base import BaseInterceptor
+
+from pythonagent.utils import get_current_timestamp_in_us, get_validated_duration, get_current_thread_time_in_us, generate_callout_id
+
+
+class PikaBlockingConnectionInterceptor(BaseInterceptor):
+
+    def __init__(self, agent, cls):
+        self.agent = agent
+        super(PikaBlockingConnectionInterceptor, self).__init__(agent, cls)
+
+    def _channel(self, channel, *args, **kwargs):
+        try:
+            connection, = args
+            select_connection = connection._impl
+            # Access host and port directly from connection
+            params = select_connection.params
+            self.host = params.host
+            self.port = params.port
+        except Exception as e:
+            self.agent.logger.warning("Exception occurred while instrumenting pika BlockingConnection.channel {}".format(e))
+
+        result = channel(*args, **kwargs)
+        return result
+
+class PikaBlockingChannelInterceptor(BaseInterceptor):
+
+    def __init__(self, agent, cls):
+        self.agent = agent
+        super(PikaBlockingChannelInterceptor, self).__init__(agent, cls)
+
+    def _basic_publish(self, basic_publish, *args, **kwargs):
+        exit_call = None
+        try:
+            channel_prop = args[0]
+            channel_impl = channel_prop._impl
+            # Access connection from channel 
+            connection = channel_impl.connection
+            connection_parameters = connection.params
+            # Access host and port from connection
+            host = connection_parameters.host
+            port = connection_parameters.port
+            db_start_time = 0
+            mNAme = basic_publish.__module__ + "." + basic_publish.__qualname__
+            db = kwargs.get('routing_key', 'default_db')
+            host_name_str = f"NA|{host}|{port}|NA|MQ|{db}|NA|NA|NA|default_user|NA"
+            ctx = self.agent.get_transaction_context()
+            if ctx:
+                ctx.entry_point_fqm = mNAme
+                ctx.host_str = host_name_str
+                ctx.status_code = 200
+            db_start_time = get_current_timestamp_in_us()
+            if ctx:
+                ctx.db_callout_start_time = db_start_time
+            db_query = None
+            db_query_params = None
+            exit_call = self.agent.db_call_begin(self.bt, host_name_str, db_query, db_query_params, db_start_time)
+
+            # Injecting custom headers
+            headers_value = self.agent.get_nd_header(self.bt, host_name_str, generate_callout_id()).decode()
+            properties = kwargs.get('properties', None)
+            if properties:
+                if properties.headers:
+                    properties.headers["CavNDFPInstance"] = headers_value
+                else:
+                    properties.headers = {}
+                    properties.headers["CavNDFPInstance"] = headers_value
+
+        except Exception as e:
+            self.agent.logger.warning("Exception in DB Begin for pika BlockingChannel.basic_publish {}".format(e))
+
+        result = basic_publish(*args, **kwargs)
+
+        try:
+            db_end_time = get_current_timestamp_in_us()
+            ctx = self.agent.get_transaction_context()
+            db_start_time = ctx.db_callout_start_time if ctx else 0
+            duration = get_validated_duration(db_start_time, db_end_time, "MQ")
+            if ctx:
+                ctx.status_code = 200
+                ctx.db_callout_start_time = 0
+            self.agent.db_call_end(self.bt, exit_call, duration, db_start_time)
+        except Exception as e:
+            self.agent.logger.warning("Exception in DB End for pika BlockingChannel.basic_publish {}".format(e))
+
+        return result
+
+    def wrapped_on_consumer_message_delivery(self, _on_consumer_message_delivery, *args, **kwargs):
+        bt = None
+        cpu_time_before = None
+        fp_instance = None
+
+        try:
+            method = args[2]
+            properties = args[3]
+
+            bt_name = method.routing_key
+            if hasattr(properties, 'headers'):
+                if isinstance(properties.headers, dict) and 'CavNDFPInstance' in properties.headers:
+                    fp_instance = properties.headers['CavNDFPInstance']
+
+            cpu_time_before = get_current_thread_time_in_us()
+            bt = self.agent.start_business_transaction(bt_name, "", None, None, None, fp_instance)
+
+        except Exception as e:
+            self.agent.logger.warning("Exception in Start BT for RabbitMQ {}".format(e))
+
+        result = _on_consumer_message_delivery(*args, **kwargs)
+
+        try:
+            if bt:
+                cpu_time_after = get_current_thread_time_in_us()
+                cpu_time = cpu_time_after - cpu_time_before
+                self.agent.end_business_transaction(bt, cpu_time)
+        except Exception as e:
+            self.agent.logger.warning("Exception in End BT for RabbitMQ {}".format(e))
+
+        return result
+
+
+def intercept_pika(agent, mod):
+    agent.logger.warning("Instrument module: pika{}".format(", mod: {}".format(mod) if mod else ""))
+    PikaBlockingConnectionInterceptor(agent, mod.adapters.blocking_connection.BlockingConnection).attach('channel')
+    pika_blocking_channel_interceptor = PikaBlockingChannelInterceptor(agent, mod.adapters.blocking_connection.BlockingChannel)
+    pika_blocking_channel_interceptor.attach('basic_publish')
+    pika_blocking_channel_interceptor.attach('_on_consumer_message_delivery',patched_method_name='wrapped_on_consumer_message_delivery')

pythonagent/agent/probes/mongodb/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import unicode_literals
+
+from .pymongo import intercept_pymongo
+
+__all__ = ['intercept_pymongo']
+