uPjGUdZddlmZmZddlmZmZmZddlm Z ddl m Z m Z m Z ddlmZddlmZddlmZdd lmZdd lmZeej2j4hZeed <eej:j4ej<j4ej2j4hZeed <d ee zde fdZ!d e de"fdZ#de$e d e de$e%fdZ&dede fdZ'de$e%d e de(e$e%e$e%ffdZ)de$e%de fdZ*dede$e fdZ+dede,e fdZ-de$e%de(e$e,e.e dzffdZ/deee e0fde(e.dffdZ1 d2de$e2de$e2d e2de$e2fd!Z3d"e$e%d#e$e%d$e$e2d%e.d&e d'e,e.d(e$e,e.de(e,e.e%e.e%fffd)Z4d*e,e.d(e$e,e.de,e.fd+Z5ej:d,ddddfde$e%d&e d ee zd-e.d%e.dzd.e dzd/e%e efdzd0e dzdefd1Z6y)3u Main compress() function — normalizes input messages, orchestrates BM25/embedding scoring, message stubbing, and retrieval tool injection. )MappingSequence)AnyFinalcast) DualCache) extract_key stub_messagetruncate_message)build_retrieval_tool)bm25_score_messages) token_counter)CompressedResult) CallTypes_ANTHROPIC_CALL_TYPES_SUPPORTED_CALL_TYPES call_typereturnc>t|tr |jS|S)zAReturn the string value for a ``CallTypes`` enum or a raw string.) isinstancervaluers a/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/compression/compress.py_normalize_call_typer$s)Y' c|tvSN)rrs r_is_anthropic_call_typer+s - --rkeysc|sgSt|g}t|s|Sddlm}|j |\}}t t t|S)z Build retrieval tool definitions in the target request schema. - Chat-completions call types: keep OpenAI function-tool schema. - Anthropic messages call type: remap to Anthropic's custom tool schema. r)AnthropicConfig)r r*litellm.llms.anthropic.chat.transformationr! _map_toolsrlistdict)rr openai_toolsr!anthropic_tools _mcp_serverss r_build_retrieval_toolsr)/sV  /56L "9 -K$3$5$@$@$N!O\ T O ,,rcontentcg}|g}|r|j}t|tr|j|nt|tr"t |D]}|j|nwt|t rg|jd}|dk(r+|jt|jddn&|dk(r!|j|jdd|rdj|S)a Convert OpenAI/Anthropic message content blocks to plain text. Text extraction policy: - Include text-bearing fields only (`text` blocks + string values). - For `tool_result`, expand into nested `content` items. - Ignore non-textual blocks (images/documents/tool metadata/thinking metadata). Implemented iteratively (stack-based) to avoid unbounded recursion. typetext tool_resultr* ) poprstrappendr$reversedr%getjoin)r*partsstackitemelement item_types r_content_to_textr<Es!E&iE yy{ dC LL  d ##D> & W% & d #(IF" S&"!567m+ TXXi45  88E?rmessagesc |tvrtd|dttd|Dcgc] }t|}}g}|D]2}|j i|dt |j ddi4||fScc}w)z Normalize each original message to a text-surrogate content for scoring. Returns: (normalized_messages, original_messages_copy) zUnsupported call_type=z# for compression. Expected one of: .r*r.)r ValueErrorsortedr%r3r<r5)r=rmoriginal_messagesnormalized_messagesmsgs r#_normalize_messages_for_compressionrFcs--$YM1TU[\qUrTsst u  HP5P!d1g5P5P-/  ""  +CGGIr,BC     1 116QsA=ct|D]3}|jddk(st|jddcSy)z1Return the text content of the last user message.roleuserr*r.)r4r5r<)r=rEs r_extract_last_user_messagerJsA!< 776?f $#CGGIr$:; ;< rct|tsgSg}|D]^}t|ts|jddk7r)|jd}t|tsK|sN|j |`|S)Nr,tool_useid)rr$r%r5r2r3)r* tool_use_idspart tool_use_ids r_extract_tool_use_idsrQsu gt $ %'L-$%  88F z ) hhtn k3 'K    ,- rct|ts tSt}|D]^}t|ts|j ddk7r)|j d}t|t sK|sN|j |`|S)Nr,r/rP)rr$setr%r5r2add)r*tool_result_idsrOrPs r_extract_tool_result_idsrVsz gt $u '*uO-$%  88F } , hh}- k3 'K    ,- rNcg}d}|t|kr||}|jddk7r|dz }-t|jd}|s|dz }O|dzt|k\rgdfS||dz}|jddk7rgdfSt|jd}|sgdfS|D] }||vsgdfcS|j ||dzh|dz }|t|kr|d fS) z Return atomic 2-message spans for Anthropic tool exchanges. Each assistant message containing `tool_use` must be immediately followed by a user message containing matching `tool_result` blocks for all tool_use ids. rrH assistantr*invalid_anthropic_tool_sequencerIN)lenr5rQrVr3)r=spansicurrentrNnext_msgrUrPs r&_extract_anthropic_tool_exchange_spansras&$&E A c(m 1+ ;;v + - FA ,W[[-CD  FA  q5CM !88 8AE? << 6 )88 828<< 3JK88 8' =K/1<<< =  aQZ  Q7 c(m : $;r.ctdt|D}tdt|Ddd}tdt|Ddd}||z|zS)ax Return indices of messages that must never be compressed: - All system messages - The last user message - The last assistant message The last user message is what the model is being asked to act on right now, so compressing it replaces the live instruction with a marker. Compression guardrails share this policy; see the Headroom guardrail. c3RK|]\}}|jdddk(s|!yw)rHr.systemNr5.0indexrEs r z(get_protected_indices..s,!qJE3QTQXQXY_acQdhpQp%!q''c3RK|]\}}|jdddk(s|!yw)rHr.rINrerfs rriz(get_protected_indices..s)jzucCGGTZ\^L_ciLiUjrjNc3RK|]\}}|jdddk(s|!yw)rHr.rXNrerfs rriz(get_protected_indices..s)mZUC#''RXZ\J]alJl5mrj)tuple enumerate)r=system_indices last_userlast_assistants rget_protected_indicesrsse"!q)H:M!qqNjYx5HjjkmknoIm9X3FmmnpnqrN I % 66r bm25_scores emb_scores bm25_weightcdttdttfd}||}||}d|z }t||Dcgc]\}}||z||zzc}}Scc}}w)zJWeighted average of BM25 and embedding scores, with min-max normalization.scoresrc|r t|nd}|r t|nd}||z }|dk(rdgt|zS|Dcgc] }||z |z  c}Scc}w)Nr)minmaxr\)rxmin_smax_srngss r _normalizez#_combine_scores.._normalizesZ&,s6{#&,s6{#U] !853v;& &+12aUc!222sAg?)r$floatzip) rtrurvr norm_bm25norm_emb emb_weightbes r_combine_scoresrsi34;34;3"+.I ,Hk)J9.s;ssl";c3(K|] }| ywrr)rgrrs rriz2_select_kept_indices_for_budget..sF#-FsFc |dS)Nrr)r9s rz1_select_kept_indices_for_budget.."s $q'r)keyreversec3&K|]}|v ywrrrs rriz2_select_kept_indices_for_budget..%s6ssl"6rrYd)rSrrr2r5roranger\r3rnrAanyr|sortupdater rT)rDrCrrrrrcurrent_tokensr^truncated_overridesspan_id_by_indexspan_idspanrcandidate_units span_indices span_score_scoreindices can_truncate msg_tokens remaining truncatedtruncated_tokensrs ` @r_select_kept_indices_for_budgetrs23LN  -c.q155iDJK   35.0"#67,  ,C$+ S ! ,, IKOS,-.E " "c\&9  4sfdCDE $BVD\* ;l; ; FFF  L%@A B14@)8/% 6g6 6   C -#237;;IrJPbQ J  '7 > $ , ,,! J &*< <    ( j (N c'la/I4D!*C():3)?KI,]]9b17R  (1  $   S ! . .N9/< , ,,rrcvt}|D](}tfd|Dr|j|*|S)Nc3&K|]}|v ywrrrs rriz1_get_dropped_tool_span_indices..H733,&7r)rSrr)rrdropped_tool_span_indicesrs` r_get_dropped_tool_span_indicesrEs>14#37$77 % , ,T 23 %$ri@ compression_triggerembedding_modelembedding_model_paramscompression_cachec n t|}t||\} } ||dzdz}t|ttt | } | |krt | | | digdSt| } t| | } |r"dd l m }|| | ||| }t| |d }n| }t| }t| g}t|rOt| \}}|t | | | dig|S|D](}t! fd |Ds j#|*t%| | ||| |\ }g}i}t}t' |}t)| D]|\}}||vr | vr"|j+|j-||1t/| |||}t1|j-dd}|||<|j+t3||~t5t |j7|}t|ttt |}t || || dkDrt9d|| z z dnd||S)u Compress a list of messages by replacing low-relevance content with stubs. Messages below ``compression_trigger`` tokens pass through unchanged. Messages above are scored with BM25 (and optionally embeddings), ranked, and the lowest-relevance messages are replaced with stubs. Originals are cached and a retrieval tool is injected so the model can recover dropped content on demand. Parameters: messages: The conversation messages to (potentially) compress. model: The LLM model name — used for token counting. call_type: The LiteLLM call type whose message schema these messages follow. Supported values: - ``CallTypes.completion`` / ``CallTypes.acompletion`` — OpenAI chat-completions shape (default) - ``CallTypes.anthropic_messages`` — Anthropic Messages shape (structured content blocks + atomic tool exchanges) compression_trigger: Only compress if input exceeds this token count. compression_target: Target token count after compression. Defaults to ``compression_trigger // 2``. embedding_model: If provided, use BM25 + embeddings for scoring. If ``None``, BM25 only. embedding_model_params: Optional kwargs forwarded to ``litellm.embedding()`` when ``embedding_model`` is set. compression_cache: Passed through to ``litellm.embedding()`` for cross-turn caching of embedding vectors. Returns: A ``CompressedResult`` dict containing compressed messages, token counts, a cache of original content, and the retrieval tool definition. )r=r )rr=rz below_trigger)r=original_tokenscompressed_tokenscompression_ratiocachetoolscompression_skipped_reasonr)embedding_score_messages)rrr皙?)rvc3&K|]}|v ywrrrs rrizcompress..rr)rDrCrrrrr)rr)fallback_index used_keysr*r.rrY)r=rrrrr)rrFrrr$rrrJr ,litellm.compression.scoring.embedding_scorerrrrsrSrrarrrrror3r5r r<r r)rround)!r=rrrrrrr call_type_strrDrCrqueryrtrrurprotected_indicesrtool_sequence_errorrrcompressed_messagesrrrr^rErr*rrrs! @rcompressrMsT0 :M-P.** !014:*d3i!23O --&+-!'6  ..ABE-U4GHK 5  !##9  *+zsS% 55HI !23L*,}-3YZk3l00  *#* /"1"%+> ( *D7$77##D) * )H/+'-)/)%L%.0#%E!$I'E!7J(-. ?3 ) )    & &':'>'>q#'F G1!4QR[\C&swwy"'=>G E#J  & &|C'= > ?*$uzz|*< VE,d3i!45  $'+RadeRe5&7/&I!JANkn  r)r)7__doc__collections.abcrrtypingrrrlitellm.caching.dual_cacher$litellm.compression.message_stubbingr r r "litellm.compression.retrieval_toolr litellm.compression.scoring.bm25r (litellm.litellm_core_utils.token_counterrlitellm.types.compressionrlitellm.types.utilsr frozensetanthropic_messagesrr__annotations__ completion acompletionrr2rboolrr$r%r)r<rnrFrJrQrSrVintraobjectrsrrrrrrrrrsi .##0 D@B6) ))*F*F*L*L)MNuN )""##$$** uIO.s.t.-c-s-tDz-,cc<24j22 4:tDz !"28d 3 49  c c#h (4j( 4C>3: %&(V7HWS&[-A$B7uSRUX7(SeSU SS %[ S,K-dK-DzK-%[K- K-  K- c( K-c#hK- 3s8T#t)_ $%K-\%S%PTUXY\U]P^%cfgjck%"+!5!5&%)"&48*.^4j^ ^3^ ^ d ^ 4Z ^!cNT1^!4'^^r