Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: audit tracing improvement #378

Open
wants to merge 13 commits into
base: main
Choose a base branch
from

Conversation

kzamlynska
Copy link
Collaborator

@kzamlynska kzamlynska commented Feb 25, 2025

closes #344

  • Review the main pathways in our library, identify areas that are not currently traced, and add tracing to those sections.
  • Sometimes we send traces that should not be included in the tracing event, such as vectors. We need to eliminate these from the traces.
  • In the CLI handler, we currently truncate inputs/outputs that are too long. We should approach this more intelligently. For example, the full prompt should always be visible
  • and formatted in a multiline style. Consider adding it as a rich output with Syntax('markdown'/'json').

Copy link
Contributor

Trivy scanning results.

@@ -88,7 +88,10 @@ def trace(self, name: str, **inputs: Any) -> Iterator[SimpleNamespace]: # noqa:
self.stop(outputs=vars(outputs), current_span=span)


def format_attributes(data: dict, prefix: str | None = None) -> dict:
# TODO: check do we need this method.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does this mean? Check how?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've checked it. I needed to run the otel tracer and see how the new formater looks like.
I removed this method.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After discussion I restored the format_attributes the way it was for otel trace handler .


def shorten_list(lst: list | tuple) -> list:
"""
Shortens a list if it's longer than 3 elements. Shortens list elements if it's long string.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Shortens a list if it's longer than 3 elements. Shortens list elements if it's long string.
Shortens a list if it's longer than opt_list_lenght elements. Shortens list elements if it's long string.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

changed

@@ -119,3 +122,112 @@ def format_attributes(data: dict, prefix: str | None = None) -> dict:
flattened[current_key] = repr(value)

return flattened


def format_attributes(data: dict, prefix: str | None = None) -> dict:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function with nested sub-functions that use local variables from the parent function pretty much looks like simulating a class with methods and instance properties. Maybe simply use a class if a class-like functionality is needed?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've changed the function to the class

Copy link
Contributor

github-actions bot commented Feb 26, 2025

badge

Code Coverage Summary

Filename                                                                                                        Stmts    Miss  Cover    Missing
------------------------------------------------------------------------------------------------------------  -------  ------  -------  --------------------------------------------
packages/__init__.py                                                                                                0       0  100.00%
packages/ragbits-cli/src/ragbits/cli/__init__.py                                                                   31       4  87.10%   73-74, 81-82
packages/ragbits-cli/src/ragbits/cli/_utils.py                                                                     23       4  82.61%   47, 65-67
packages/ragbits-cli/src/ragbits/cli/state.py                                                                      79       3  96.20%   50-51, 61
packages/ragbits-cli/tests/unit/test_state.py                                                                      72       2  97.22%   103-104
packages/ragbits-conversations/src/ragbits/conversations/__init__.py                                                0       0  100.00%
packages/ragbits-conversations/src/ragbits/conversations/history/__init__.py                                        0       0  100.00%
packages/ragbits-conversations/src/ragbits/conversations/history/compressors/__init__.py                            3       0  100.00%
packages/ragbits-conversations/src/ragbits/conversations/history/compressors/base.py                               10       0  100.00%
packages/ragbits-conversations/src/ragbits/conversations/history/compressors/llm.py                                29       1  96.55%   79
packages/ragbits-conversations/src/ragbits/conversations/history/stores/__init__.py                                 3       0  100.00%
packages/ragbits-conversations/src/ragbits/conversations/history/stores/base.py                                    17       0  100.00%
packages/ragbits-conversations/src/ragbits/conversations/history/stores/sql.py                                     66       1  98.48%   128
packages/ragbits-conversations/tests/unit/history/test_llm_compressor.py                                           64       0  100.00%
packages/ragbits-conversations/tests/unit/history/test_sql_store.py                                                48       2  95.83%   29-30
packages/ragbits-core/src/ragbits/core/__init__.py                                                                  6       2  66.67%   8-9
packages/ragbits-core/src/ragbits/core/cli.py                                                                       6       0  100.00%
packages/ragbits-core/src/ragbits/core/config.py                                                                   17       0  100.00%
packages/ragbits-core/src/ragbits/core/options.py                                                                  17       0  100.00%
packages/ragbits-core/src/ragbits/core/types.py                                                                     9       0  100.00%
packages/ragbits-core/src/ragbits/core/audit/__init__.py                                                           74       5  93.24%   42-45, 52-54
packages/ragbits-core/src/ragbits/core/audit/base.py                                                              116      20  82.76%   143, 188-191, 198-201, 257-276
packages/ragbits-core/src/ragbits/core/audit/cli.py                                                                91       3  96.70%   94-97
packages/ragbits-core/src/ragbits/core/embeddings/__init__.py                                                       5       0  100.00%
packages/ragbits-core/src/ragbits/core/embeddings/base.py                                                          20       2  90.00%   56, 69
packages/ragbits-core/src/ragbits/core/embeddings/exceptions.py                                                    17       7  58.82%   7-8, 17, 26-27, 36, 45
packages/ragbits-core/src/ragbits/core/embeddings/litellm.py                                                       38      18  52.63%   78-112
packages/ragbits-core/src/ragbits/core/embeddings/noop.py                                                           8       0  100.00%
packages/ragbits-core/src/ragbits/core/embeddings/sparse.py                                                        55      27  50.91%   23-24, 27, 66-94
packages/ragbits-core/src/ragbits/core/llms/__init__.py                                                             4       0  100.00%
packages/ragbits-core/src/ragbits/core/llms/base.py                                                                78       3  96.15%   67, 79, 256
packages/ragbits-core/src/ragbits/core/llms/exceptions.py                                                          17       7  58.82%   7-8, 17, 26-27, 36, 45
packages/ragbits-core/src/ragbits/core/llms/factory.py                                                             12       3  75.00%   30, 41, 51
packages/ragbits-core/src/ragbits/core/llms/litellm.py                                                             95      29  69.47%   94, 100-121, 165, 175, 204-229, 251-256, 267
packages/ragbits-core/src/ragbits/core/llms/local.py                                                               56      27  51.79%   9-12, 66-74, 86-87, 108-119, 140-156
packages/ragbits-core/src/ragbits/core/llms/mock.py                                                                29       3  89.66%   69-72
packages/ragbits-core/src/ragbits/core/metadata_stores/__init__.py                                                  3       0  100.00%
packages/ragbits-core/src/ragbits/core/metadata_stores/base.py                                                     11       0  100.00%
packages/ragbits-core/src/ragbits/core/metadata_stores/exceptions.py                                                4       0  100.00%
packages/ragbits-core/src/ragbits/core/metadata_stores/in_memory.py                                                17       0  100.00%
packages/ragbits-core/src/ragbits/core/prompt/__init__.py                                                           2       0  100.00%
packages/ragbits-core/src/ragbits/core/prompt/_cli.py                                                              44      21  52.27%   25-33, 47-49, 63-65, 73-75, 89-97
packages/ragbits-core/src/ragbits/core/prompt/base.py                                                              28       0  100.00%
packages/ragbits-core/src/ragbits/core/prompt/parsers.py                                                           35       0  100.00%
packages/ragbits-core/src/ragbits/core/prompt/prompt.py                                                           126       0  100.00%
packages/ragbits-core/src/ragbits/core/prompt/discovery/__init__.py                                                 2       0  100.00%
packages/ragbits-core/src/ragbits/core/prompt/discovery/prompt_discovery.py                                        33       2  93.94%   54-55
packages/ragbits-core/src/ragbits/core/utils/__init__.py                                                            0       0  100.00%
packages/ragbits-core/src/ragbits/core/utils/_pyproject.py                                                         38       1  97.37%   113
packages/ragbits-core/src/ragbits/core/utils/config_handling.py                                                    72       8  88.89%   16, 54-55, 62-63, 152-154
packages/ragbits-core/src/ragbits/core/utils/decorators.py                                                         29       0  100.00%
packages/ragbits-core/src/ragbits/core/utils/dict_transformations.py                                               72       3  95.83%   24, 27, 108
packages/ragbits-core/src/ragbits/core/vector_stores/__init__.py                                                    3       0  100.00%
packages/ragbits-core/src/ragbits/core/vector_stores/_cli.py                                                       51       4  92.16%   63, 85, 91, 121
packages/ragbits-core/src/ragbits/core/vector_stores/base.py                                                       40       0  100.00%
packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py                                                     59       1  98.31%   77
packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py                                                  38       0  100.00%
packages/ragbits-core/src/ragbits/core/vector_stores/pgvector.py                                                  124      12  90.32%   173-181, 215-217, 289
packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py                                                     67       1  98.51%   82
packages/ragbits-core/tests/cli/__init__.py                                                                         0       0  100.00%
packages/ragbits-core/tests/cli/test_cli_trace_handler.py                                                          48       3  93.75%   30, 43, 56
packages/ragbits-core/tests/cli/test_vector_store.py                                                              105       0  100.00%
packages/ragbits-core/tests/integration/vector_stores/test_vector_store.py                                         31       0  100.00%
packages/ragbits-core/tests/unit/__init__.py                                                                        0       0  100.00%
packages/ragbits-core/tests/unit/test_options.py                                                                   21       0  100.00%
packages/ragbits-core/tests/unit/audit/__init__.py                                                                  0       0  100.00%
packages/ragbits-core/tests/unit/audit/test_cli.py                                                                107       0  100.00%
packages/ragbits-core/tests/unit/audit/test_trace.py                                                               97       3  96.91%   16, 19, 22
packages/ragbits-core/tests/unit/embeddings/__init__.py                                                             0       0  100.00%
packages/ragbits-core/tests/unit/embeddings/test_from_config.py                                                    20       0  100.00%
packages/ragbits-core/tests/unit/llms/__init__.py                                                                   0       0  100.00%
packages/ragbits-core/tests/unit/llms/test_base.py                                                                 98       0  100.00%
packages/ragbits-core/tests/unit/llms/test_from_config.py                                                          16       0  100.00%
packages/ragbits-core/tests/unit/llms/test_litellm.py                                                              70       0  100.00%
packages/ragbits-core/tests/unit/llms/factory/__init__.py                                                           0       0  100.00%
packages/ragbits-core/tests/unit/llms/factory/test_get_preferred_llm.py                                            12       0  100.00%
packages/ragbits-core/tests/unit/metadata_stores/__init__.py                                                        0       0  100.00%
packages/ragbits-core/tests/unit/metadata_stores/test_from_config.py                                               11       0  100.00%
packages/ragbits-core/tests/unit/metadata_stores/test_in_memory.py                                                 22       0  100.00%
packages/ragbits-core/tests/unit/prompts/__init__.py                                                                0       0  100.00%
packages/ragbits-core/tests/unit/prompts/test_parsers.py                                                           65       0  100.00%
packages/ragbits-core/tests/unit/prompts/test_prompt.py                                                           165       0  100.00%
packages/ragbits-core/tests/unit/prompts/discovery/__init__.py                                                      0       0  100.00%
packages/ragbits-core/tests/unit/prompts/discovery/prompt_classes_for_tests.py                                     30       0  100.00%
packages/ragbits-core/tests/unit/prompts/discovery/test_prompt_discovery.py                                        18       0  100.00%
packages/ragbits-core/tests/unit/prompts/discovery/ragbits_tests_pkg_with_prompts/__init__.py                       2       1  50.00%   3
packages/ragbits-core/tests/unit/prompts/discovery/ragbits_tests_pkg_with_prompts/prompts/__init__.py               3       2  33.33%   2-4
packages/ragbits-core/tests/unit/prompts/discovery/ragbits_tests_pkg_with_prompts/prompts/temp_prompt1.py          14       0  100.00%
packages/ragbits-core/tests/unit/prompts/discovery/ragbits_tests_pkg_with_prompts/prompts/temp_prompt2.py          14       0  100.00%
packages/ragbits-core/tests/unit/utils/__init__.py                                                                  0       0  100.00%
packages/ragbits-core/tests/unit/utils/test_config_handling.py                                                     64       2  96.88%   27-28
packages/ragbits-core/tests/unit/utils/test_decorators.py                                                          26       2  92.31%   17, 39
packages/ragbits-core/tests/unit/utils/test_dict_transformations.py                                                69       0  100.00%
packages/ragbits-core/tests/unit/utils/pyproject/test_find.py                                                      13       0  100.00%
packages/ragbits-core/tests/unit/utils/pyproject/test_get_config.py                                                 9       0  100.00%
packages/ragbits-core/tests/unit/utils/pyproject/test_get_instace.py                                               37       0  100.00%
packages/ragbits-core/tests/unit/vector_stores/__init__.py                                                          0       0  100.00%
packages/ragbits-core/tests/unit/vector_stores/test_chroma.py                                                      61       0  100.00%
packages/ragbits-core/tests/unit/vector_stores/test_from_config.py                                                 42       0  100.00%
packages/ragbits-core/tests/unit/vector_stores/test_in_memory.py                                                   84       0  100.00%
packages/ragbits-core/tests/unit/vector_stores/test_pgvector.py                                                   173       0  100.00%
packages/ragbits-core/tests/unit/vector_stores/test_qdrant.py                                                      64       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/__init__.py                                            2       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/_main.py                                             119       4  96.64%   266-273, 276
packages/ragbits-document-search/src/ragbits/document_search/cli.py                                                39       2  94.87%   85, 104
packages/ragbits-document-search/src/ragbits/document_search/documents/__init__.py                                  0       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/documents/document.py                                 62       2  96.77%   100, 147
packages/ragbits-document-search/src/ragbits/document_search/documents/element.py                                  78      12  84.62%   87, 167-174, 183-185
packages/ragbits-document-search/src/ragbits/document_search/documents/exceptions.py                               11       1  90.91%   17
packages/ragbits-document-search/src/ragbits/document_search/documents/sources/__init__.py                          7       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/documents/sources/azure.py                            91      11  87.91%   71-72, 106-107, 177-186
packages/ragbits-document-search/src/ragbits/document_search/documents/sources/base.py                             64       4  93.75%   166-167, 170-171
packages/ragbits-document-search/src/ragbits/document_search/documents/sources/gcs.py                              58       1  98.28%   44
packages/ragbits-document-search/src/ragbits/document_search/documents/sources/hf.py                               54      12  77.78%   52-55, 59-60, 90, 97-98, 112-114
packages/ragbits-document-search/src/ragbits/document_search/documents/sources/local.py                            36       2  94.44%   38, 76
packages/ragbits-document-search/src/ragbits/document_search/documents/sources/s3.py                               99      45  54.55%   49-56, 70-94, 114-129, 160, 177
packages/ragbits-document-search/src/ragbits/document_search/ingestion/__init__.py                                  0       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/document_processor.py                       33       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/processor_strategies/__init__.py             5       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/processor_strategies/base.py                25       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/processor_strategies/batched.py             18       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/processor_strategies/distributed.py         30       8  73.33%   8-9, 36, 64-71
packages/ragbits-document-search/src/ragbits/document_search/ingestion/processor_strategies/sequential.py          13       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/__init__.py                        3       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/base.py                           19       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/dummy.py                          20       7  65.00%   33, 54-60
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/__init__.py           4       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/default.py           46       4  91.30%   98, 103-104, 137
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/images.py            50      19  62.00%   73-80, 87-99, 111, 124
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/pdf.py               19       6  68.42%   23, 35-43
packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/utils.py             38      11  71.05%   71, 82-83, 98-101, 110, 121-123
packages/ragbits-document-search/src/ragbits/document_search/retrieval/__init__.py                                  0       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/__init__.py                       6       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/base.py                           9       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/llm.py                           25       4  84.00%   47-50
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/multi.py                         27       4  85.19%   51-54
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/noop.py                           6       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/prompts.py                       26       2  92.31%   65, 87
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/__init__.py                        3       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/base.py                           17       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/litellm.py                        18       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/noop.py                           11       0  100.00%
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/reciprocal_ranked_fusion.py       22       2  90.91%   50, 60
packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/rerankers_answerdotai.py          21       0  100.00%
packages/ragbits-document-search/tests/__init__.py                                                                  0       0  100.00%
packages/ragbits-document-search/tests/helpers.py                                                                   3       0  100.00%
packages/ragbits-document-search/tests/cli/test_ingest.py                                                          21       0  100.00%
packages/ragbits-document-search/tests/cli/test_search.py                                                          71       0  100.00%
packages/ragbits-document-search/tests/integration/__init__.py                                                      0       0  100.00%
packages/ragbits-document-search/tests/integration/test_rerankers.py                                               24       6  75.00%   20-42
packages/ragbits-document-search/tests/integration/test_sources.py                                                 24      10  58.33%   23-33, 41-46
packages/ragbits-document-search/tests/integration/test_unstructured.py                                            48      10  79.17%   52-58, 71-77
packages/ragbits-document-search/tests/unit/test_aws_source.py                                                     24       0  100.00%
packages/ragbits-document-search/tests/unit/test_azure_blob_source.py                                              76       0  100.00%
packages/ragbits-document-search/tests/unit/test_config.py                                                         57       0  100.00%
packages/ragbits-document-search/tests/unit/test_document_processor.py                                             17       0  100.00%
packages/ragbits-document-search/tests/unit/test_document_search.py                                               233       1  99.57%   463
packages/ragbits-document-search/tests/unit/test_documents.py                                                      13       0  100.00%
packages/ragbits-document-search/tests/unit/test_elements.py                                                       19       0  100.00%
packages/ragbits-document-search/tests/unit/test_gcs_hf_sources.py                                                 53       8  84.91%   17-18, 53-58
packages/ragbits-document-search/tests/unit/test_local_file_source.py                                              13       0  100.00%
packages/ragbits-document-search/tests/unit/test_processing_strategies.py                                          25       0  100.00%
packages/ragbits-document-search/tests/unit/test_providers.py                                                      41       0  100.00%
packages/ragbits-document-search/tests/unit/test_rephrasers.py                                                     40       0  100.00%
packages/ragbits-document-search/tests/unit/test_rerankers.py                                                      81       1  98.77%   25
packages/ragbits-document-search/tests/unit/test_source_discriminator.py                                           36       0  100.00%
packages/ragbits-document-search/tests/unit/testprojects/project_with_instance_factory/__init__.py                  0       0  100.00%
packages/ragbits-document-search/tests/unit/testprojects/project_with_instance_factory/factories.py                20       0  100.00%
packages/ragbits-evaluate/src/ragbits/evaluate/__init__.py                                                          0       0  100.00%
packages/ragbits-evaluate/src/ragbits/evaluate/cli.py                                                              45      16  64.44%   109-119, 130-149
packages/ragbits-evaluate/src/ragbits/evaluate/config.py                                                            6       0  100.00%
packages/ragbits-evaluate/src/ragbits/evaluate/evaluator.py                                                        51      25  50.98%   43-48, 71-78, 99-109, 122, 136, 151-155
packages/ragbits-evaluate/src/ragbits/evaluate/dataloaders/__init__.py                                              7       2  71.43%   20-21
packages/ragbits-evaluate/src/ragbits/evaluate/dataloaders/base.py                                                  7       0  100.00%
packages/ragbits-evaluate/src/ragbits/evaluate/metrics/__init__.py                                                  2       0  100.00%
packages/ragbits-evaluate/src/ragbits/evaluate/metrics/base.py                                                     21       5  76.19%   24-25, 54, 67, 79
packages/ragbits-evaluate/src/ragbits/evaluate/pipelines/__init__.py                                               11       4  63.64%   23-26
packages/ragbits-evaluate/src/ragbits/evaluate/pipelines/base.py                                                   14       2  85.71%   23, 41
packages/ragbits-evaluate/src/ragbits/evaluate/pipelines/document_search.py                                        29      12  58.62%   35-36, 52-55, 61-68, 80-82
packages/ragbits-guardrails/src/ragbits/guardrails/__init__.py                                                      0       0  100.00%
packages/ragbits-guardrails/src/ragbits/guardrails/base.py                                                         15       0  100.00%
packages/ragbits-guardrails/src/ragbits/guardrails/openai_moderation.py                                            19       5  73.68%   29-33
packages/ragbits-guardrails/tests/unit/test_openai_moderation.py                                                   35       0  100.00%
TOTAL                                                                                                            6159     509  91.74%

Diff against main

Filename                                                Stmts    Miss  Cover
----------------------------------------------------  -------  ------  --------
packages/ragbits-core/src/ragbits/core/audit/base.py      +72     +19  -14.97%
packages/ragbits-core/src/ragbits/core/audit/cli.py       +22      +3  -3.30%
packages/ragbits-core/tests/unit/audit/test_cli.py         +5       0  +100.00%
packages/ragbits-core/tests/unit/audit/test_trace.py       +4       0  +0.14%
TOTAL                                                    +103     +22  -0.22%

Results for commit: 99800af

Minimum allowed coverage is 60%

♻️ This comment has been updated with latest results

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

feat: audit tracing improvements
2 participants