contextflow-ai 0.1.2__py3-none-any.whl

contextflow/core/compactor.py

@@ -0,0 +1,120 @@
+"""
+Implements message summarization techniques
+"""
+
+from typing import List, Dict
+from utils.llm import LLMClient
+
+
+class MessageCompactor:
+    def __init__(self):
+        self.llm = LLMClient()
+
+    def summarize(
+        self,
+        messages_to_summarize: List[Dict[str, str]],
+        max_token_count: int = 500,
+    ) -> str:
+        """
+        Summarizes a list of messages.
+
+        Args:
+            messages_to_summarize: The list of messages to compress
+            agent_goal: The overall goal to guide the summary
+            max_token_count: The target length for the final summary
+
+        Returns:
+            summaries: A single string containing the dense summary.
+        """
+        return self._simple_summarize(messages_to_summarize, max_token_count)
+
+    def _simple_summarize(
+        self,
+        messages_to_summarize: List[Dict[str, str]],
+        max_token_count: int,
+    ):
+        """
+        Summarizes a list of messages using a simple LLM summary.
+
+        Args:
+            messages_to_summarize: The list of messages to compress
+            agent_goal: The overall goal to guide the summary
+            max_token_count: The target length for the final summary
+
+        Returns:
+            summaries: A single string containing the dense summary.
+        """
+        if not messages_to_summarize:
+            return ""
+
+        if len(messages_to_summarize) == 1:
+            return messages_to_summarize[0]["content"]
+
+        conversation_text = self._format_messages(messages_to_summarize)
+
+        prompt = f"""You are summarizing a conversation to preserve key information while reducing length.
+
+        Conversation:
+        {conversation_text}
+
+        Instructions:
+        - Create a dense, information-rich summary
+        - Preserve all critical facts, names, numbers, and decisions
+        - Remove pleasantries and redundant information
+        - Target length: approximately {max_token_count} tokens
+        - Write in third person (e.g., "User reported X. Agent confirmed Y.")
+
+        Summary:"""
+
+        try:
+            summary = self.llm.generate_text(
+                prompt,
+                max_tokens=max_token_count,
+                temperature=0.3,  # Low temperature for factual summary
+            )
+
+            return summary.strip()
+        except Exception as e:
+            # Fallback: return a simple concatenation
+            print(f"Warning: Summarization failed ({e}). Using fallback.")
+            return self._fallback_summary(messages_to_summarize)
+
+    def _format_messages(self, messages: List[Dict[str, str]]) -> str:
+        """Format messages into a readable conversation"""
+
+        formatted = []
+        for msg in messages:
+            role = msg.get("role", "unknown")
+            content = msg.get("content", "")
+            formatted.append(f"{role.capitalize()}: {content}")
+
+        return "\n".join(formatted)
+
+    def _fallback_summary(self, messages: List[Dict[str, str]]) -> str:
+        """Simple fallback if LLM summarization fails"""
+
+        # Just concatenate the messages with "..." between them
+        contents = [msg.get("content", "") for msg in messages]
+        return " ... ".join(contents)
+
+    def _hierarchical_summarize(
+        self,
+        messages_to_summarize: List[Dict[str, str]],
+        agent_goal: str,
+        max_token_count: int,
+    ):
+        """
+        Summarizes a list of messages using the map-reduce method.
+
+        Args:
+            messages_to_summarize: The list of messages to compress
+            agent_goal: The overall goal to guide the summary
+            max_summary_tokens: The target length for the final summary
+
+        Returns:
+            summaries: A single string containing the dense summary.
+        """
+        raise NotImplementedError(
+            "Hierarchical summarization not yet implemented. "
+            "Use simple_summarize for now."
+        )

contextflow/core/optimizer.py

@@ -0,0 +1,39 @@
+from core.compactor import MessageCompactor
+from core.scorer import MessageScorer
+from typing import List, Dict, Union, Literal
+from core.strategies import balanced_strategy
+from utils.tokenizer import count_tokens
+
+
+class ContextFlow:
+    def __init__(self):
+        self.message_compactor = MessageCompactor()
+        self.message_scorer = MessageScorer()
+
+    def optimize(
+        self,
+        messages: List[Dict[str, str]],
+        agent_goal: str,
+        max_token_count: int = 500,
+    ):
+        scores = self.message_scorer.score_messages(
+            messages=messages, agent_goal=agent_goal
+        )
+
+        optimized = None
+        optimized = balanced_strategy(
+            messages, scores, max_token_count, self.message_compactor
+        )
+
+        tokens_before = count_tokens(messages)
+        tokens_after = count_tokens(optimized)
+        reduction_pct = ((tokens_before - tokens_after) / tokens_before) * 100
+
+        return {
+            "messages": optimized,
+            "analytics": {
+                "tokens_after": tokens_after,
+                "reduction_pct": reduction_pct,
+                "tokens_saved": tokens_before - tokens_after,
+            },
+        }

contextflow/core/scorer.py

@@ -0,0 +1,82 @@
+"""
+Message relevance and utility scoring
+"""
+
+from typing import List, Dict
+from utils.llm import LLMClient
+import numpy as np
+
+
+class MessageScorer:
+    def __init__(self):
+        self._embedding_cache = {}
+        self._llm_client = LLMClient()
+
+    def _encode_text(self, text: str) -> np.ndarray:
+        """Convert text to a vector embedding
+
+        Args:
+            text (str): Text to encode
+        Returns:
+            A 384-dimensional numpy array representing the text encoding
+        """
+
+        if text in self._embedding_cache:
+            return self._embedding_cache[text]
+
+        embedding = self.embedding_model.encode(
+            text,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+
+        self._embedding_cache[text] = embedding
+
+        return embedding
+
+    def _create_batches(self, messages: List[Dict[str, str]]):
+        if len(messages) < 20:
+            return [messages]
+
+        return [messages[i : i + 20] for i in range(0, len(messages), 20)]
+
+    def score_messages(
+        self, messages: List[Dict[str, str]], agent_goal: str
+    ) -> List[float]:
+        """Scores messages based on how relevant they are to the agent's goal.
+
+        Args:
+            messages: A list of messages
+            agent_goal: The goal of the agent
+        Returns:
+            A list scores such that scores[i] is the relevancy score of messages[i]
+        """
+
+        scores = []
+
+        batches = self._create_batches(messages)
+
+        for i, batch in enumerate(batches):
+            batch_scores = self._llm_client.batch_score_utility(
+                batch, agent_goal
+            )
+            scores += batch_scores
+
+        for i in range(len(scores) - 1, max(-1, len(scores) - 6), -1):
+            recency_bonus = 1.0
+            scores[i] += recency_bonus
+
+        return scores
+
+    def _calculate_utility(self, text: str, goal: str) -> float:
+        """Returns a utility score (0-10 scale)
+
+        Args:
+            text: String to evaluate the utility of
+        Returns:
+            Utility score of `text` between 0 and 10
+        """
+        text = text.strip()
+        text_lower = text.lower()
+
+        return self._llm_client.score_utility(text_lower, goal)

contextflow/core/strategies.py

@@ -0,0 +1,128 @@
+from typing import List, Union, Literal
+from utils.tokenizer import count_tokens  # You'll need this
+from core.compactor import MessageCompactor
+
+
+def get_strategy(
+    strategy_name: Union[
+        Literal["conservative"], Literal["balanced"], Literal["aggressive"]
+    ],
+):
+    """Factory to get strategy function"""
+    strategies = {
+        "conservative": conservative_strategy,
+        "balanced": balanced_strategy,
+        "aggressive": aggressive_strategy,
+    }
+    return strategies.get(strategy_name, balanced_strategy)
+
+
+def conservative_strategy():
+    pass
+
+
+def balanced_strategy(
+    messages: List[str],
+    scores: List[float],
+    max_token_count: int,
+    compactor: MessageCompactor,
+):
+    """Optimizes a conversation (i.e. a list of messages) using a balanced strategy (keep high-scoring, summarize mid, drop low)
+
+    Args:
+        messages: List of messages
+        scores: List of scores for each message
+        max_token_count: Maximum number of tokens allowed
+        compactor: Tool for summarizing messages
+    Returns:
+        Optimized list of messages that is less than max_token_count
+    """
+
+    preserve_recent = 5
+
+    recent_scores = scores[-10:] if len(scores) >= 10 else scores
+
+    # If the recent messages are high-utility, keep more
+    avg_recent_score = sum(recent_scores) / len(recent_scores)
+
+    if avg_recent_score >= 7:
+        preserve_recent = 5  # Keep 5 if they're useful
+    elif avg_recent_score >= 4:
+        preserve_recent = 3  # Keep 3 if they're medium
+    else:
+        preserve_recent = 2  # Keep only 2 if they're low-utility pleasantries
+
+    print(f"preserving: {preserve_recent}")
+
+    recent_messages = messages[-preserve_recent:]
+    older_messages = messages[:-preserve_recent]
+    older_scores = scores[:-preserve_recent]
+
+    sorted_pairs = sorted(
+        zip(older_messages, older_scores),
+        key=lambda x: x[1],  # Sort by the score (second element of tuple)
+        reverse=True,  # Highest scores first
+    )
+
+    optimized = recent_messages.copy()
+    current_tokens = count_tokens(optimized)
+
+    # FIX: Check if we're already over budget with just recent messages
+    if current_tokens >= max_token_count:
+        # Emergency: Even recent messages exceed budget
+        # Keep reducing until we fit
+        while current_tokens > max_token_count and len(optimized) > 1:
+            optimized.pop(0)  # Remove oldest of the recent messages
+            current_tokens = count_tokens(optimized)
+        return optimized
+
+    # Categorize older messages into buckets
+    keep_bucket = []
+    summarize_bucket = []
+
+    for message, score in sorted_pairs:
+        if score > 7.0:
+            keep_bucket.append(message)
+        elif score > 4.0:
+            summarize_bucket.append(message)
+        # score <= 4.0: drop entirely
+
+    # Try to add high-scoring messages one by one
+    for message in keep_bucket:
+        message_tokens = count_tokens([message])
+        if current_tokens + message_tokens <= max_token_count:
+            optimized.insert(0, message)  # Add before recent messages
+            current_tokens += message_tokens
+        else:
+            # print(f"adding {message} to summary...")
+            # Can't fit this message, add to summarize bucket instead
+            summarize_bucket.append(message)
+
+    # print(summarize_bucket)
+    if summarize_bucket:
+        summary_tokens = count_tokens(summarize_bucket)
+        print(f"Tokens needed to summarize: {summary_tokens}")
+        summary = compactor.summarize(summarize_bucket, summary_tokens * 0.2)
+        summary_message = {
+            "role": "system",
+            "content": f"Summary of earlier context: {summary}",
+        }
+        print(f"Tokens summary: {count_tokens([summary_message])}")
+
+        # print(summary_message)
+
+        summary_tokens = count_tokens([summary_message])
+        if current_tokens + summary_tokens <= max_token_count:
+            optimized.insert(0, summary_message)
+            current_tokens += summary_tokens
+        else:
+            print(
+                f"Summary tokens exceeds {max_token_count}. Dropping summary."
+            )
+        # If summary doesn't fit, skip it (rare but possible)
+
+    return optimized
+
+
+def aggressive_strategy():
+    pass

contextflow/utils/llm.py

@@ -0,0 +1,206 @@
+"""
+Local LLM client using Ollama
+"""
+
+from google import genai
+from google.genai import types
+import json
+import os
+from typing import List, Dict
+
+
+class LLMClient:
+    """Lightweight LLM utility"""
+
+    def __init__(self):
+        self.client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+        self.model_name = "gemini-2.5-flash-lite"
+
+    def generate_text(
+        self, prompt: str, max_tokens: int = 500, temperature: float = 0.3
+    ) -> str:
+        """
+        Generate text from a prompt (for summarization)
+
+        Args:
+            prompt: The prompt to send
+            max_tokens: Maximum tokens in response
+            temperature: Randomness (0=deterministic, 1=creative)
+
+        Returns:
+            Generated text
+        """
+
+        try:
+            response = self.client.models.generate_content(
+                model=self.model_name,
+                contents=prompt,
+                config=types.GenerateContentConfig(
+                    temperature=0,
+                ),
+            )
+
+            return response.text
+
+        except Exception as e:
+            raise Exception(f"LLM generation failed: {e}")
+
+    def score_utility(self, text: str, goal: str) -> float:
+        """
+        Score message utility using an LLM
+
+        Args:
+            text: Message content
+            goal: Agent goal for context
+
+        Returns:
+            Utility score 0-10
+        """
+
+        prompt = f"""Rate message relevance to goal (0-10 scale):
+
+        Goal: {goal}
+
+        Scoring guide:
+        • 8-10: Specific facts, numbers, IDs, errors ("Error: timeout line 42" = 10)
+        • 4-7: Questions, partial info ("Can you check status?" = 6)
+        • 0-3: Acknowledgments, greetings ("Thanks!" = 1)
+
+        Message: "{text}"
+
+        Score (number only):"""
+
+        response = self.client.models.generate_content(
+            model=self.model_name,
+            contents=prompt,
+            config=types.GenerateContentConfig(
+                temperature=0,
+                response_mime_type="application/json",
+                response_schema={
+                    "type": "object",
+                    "properties": {
+                        "score": {
+                            "type": "number",
+                            "description": "Rating from 0-10",
+                            "minimum": 0,
+                            "maximum": 10,
+                        }
+                    },
+                    "required": ["score"],
+                },
+            ),
+        )
+        result = json.loads(response.text)
+        score = result["score"]
+
+        return score
+
+    def batch_score_utility(
+        self, batch: List[Dict[str, str]], goal: str
+    ) -> List[float]:
+        """
+        Score message utility using an LLM
+
+        Args:
+            text: Message content
+            goal: Agent goal for context
+
+        Returns:
+            Utility score 0-10
+        """
+
+        formatted_messages = ""
+        for i, msg in enumerate(batch, 1):
+            role = msg.get("role", "unknown").capitalize()
+            content = msg.get("content", "")
+            # Truncate long messages to save prompt tokens?
+            # formatted_messages += f"{i}. [{role}] {content[:200]}\n"
+            formatted_messages += f"{i}. [{role}] {content}\n"
+
+        prompt = f"""Rate message relevance to goal (0-10 scale):
+
+        Goal: {goal}
+
+        Scoring guide:
+        • 8-10: Specific facts, numbers, IDs, errors ("Error: timeout line 42" = 10)
+        • 4-7: Questions, partial info ("Can you check status?" = 6)
+        • 0-3: Acknowledgments, greetings ("Thanks!" = 1)
+
+        MESSAGES TO RATE:
+        {formatted_messages}
+
+        Return a JSON array with one score per message in order
+        """
+        response_schema = {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "message_index": {
+                        "type": "integer",
+                        "description": "The message number (1-indexed)",
+                    },
+                    "score": {
+                        "type": "number",
+                        "description": "Utility score from 0-10",
+                        "minimum": 0,
+                        "maximum": 10,
+                    },
+                },
+                "required": ["message_index", "score"],
+            },
+        }
+
+        try:
+            response = self.client.models.generate_content(
+                model=self.model_name,
+                contents=prompt,
+                config=types.GenerateContentConfig(
+                    temperature=0,
+                    response_mime_type="application/json",
+                    response_schema=response_schema,
+                ),
+            )
+
+            scores_data = json.loads(response.text)
+
+            scores = self._extract_scores_from_json(scores_data, len(batch))
+
+            return scores
+
+        except Exception as e:
+            print(
+                f"Warning: Batch scoring failed ({e}). Using fallback scores."
+            )
+            return [5.0] * len(batch)
+
+    def _extract_scores_from_json(
+        self, scores_data: List[Dict], expected_count: int
+    ) -> List[float]:
+        """
+        Extracts scores from the JSON array response.
+
+        Args:
+            scores_data: List of {"message_index": 1, "score": 8.5} objects
+            expected_count: Number of messages we scored
+
+        Returns:
+            List of scores in the correct order
+        """
+
+        # Create a dict for quick lookup
+        score_dict = {}
+        for item in scores_data:
+            idx = item.get("message_index")
+            score = item.get("score", 5.0)
+            if idx is not None:
+                score_dict[idx] = max(0.0, min(10.0, float(score)))
+
+        # Build the final list in order (1-indexed to 0-indexed)
+        final_scores = []
+        for i in range(1, expected_count + 1):
+            final_scores.append(
+                score_dict.get(i, 5.0)
+            )  # Default to 5.0 if missing
+
+        return final_scores

contextflow/utils/tokenizer.py

@@ -0,0 +1,26 @@
+"""
+Token counting utilities
+"""
+
+from typing import List, Dict
+
+
+def count_tokens(messages: List[Dict[str, str]]) -> int:
+    """
+    Count tokens in a list of messages
+
+    Args:
+        messages: List of message dicts [{"role": "user", "content": "..."}]
+
+    Returns:
+        Total token count
+    """
+
+    # Hueristic. Not exact
+
+    total = 0
+
+    for message in messages:
+        total += len(message["content"])
+
+    return total // 4

contextflow_ai-0.1.2.dist-info/METADATA

@@ -0,0 +1,77 @@
+Metadata-Version: 2.4
+Name: contextflow-ai
+Version: 0.1.2
+Summary:  Context optimization for AI agents
+Author-email: Ayush Thapa <ayushthapa026@gmail.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: google-genai==1.49.0
+Requires-Dist: click==8.1.7
+Requires-Dist: rich==13.9.4
+Requires-Dist: python-dotenv==1.0.1
+Requires-Dist: numpy==2.3.4
+Dynamic: license-file
+
+<p align="center">
+  <p align="center">
+	<img width="128" height="128" src="https://github.com/athaapa/contextflow/blob/74d4dffdb3c0c1777ccd3593b7d52f89f9406afc/images/logo.png" alt="Logo">
+  </p>
+
+  <h1 align="center"><b>ContextFlow</b></h1>
+  <p align="center">
+     Context optimization for AI agents.
+    <br />
+  </p>
+</p>
+
+**ContextFlow** is a fast, open-source Python library that compresses chat, agent, or RAG context before you send it to an expensive LLM. It finds the “needle in the haystack,” scoring every message by utility, keeping only critical facts, summarizing the rest, and dropping noise—preserving quality while saving money.
+
+# Why ContextFlow?
+Most LLM applications waste money by sending entire chat histories, long chains of document chunks, or endless boilerplate every API call. Most of it is redundant, repetitive, or irrelevant to the user’s actual goal.
+
+ContextFlow solves this by:
+- Ranking every message by utility using a fast LLM batch
+- Keeping only high-signal content (order numbers, errors, decisions)
+- Aggressively summarizing medium-utility content
+- Dropping low-value fluff and polite filler
+- Always preserving the last few recent messages for context
+
+# Quickstart
+## Installation
+```bash
+pip install contextflow
+```
+## Setup
+Get your Gemini API Key (for message scoring and summarization):
+```.env
+export GEMINI_API_KEY="YOUR_KEY_HERE"
+```
+If you don't have one, you can get one for free [here](https://aistudio.google.com/api-keys).
+## Example
+```python
+from contextflow import ContextFlow
+
+# Chat history (list of {"role": str, "content": str})
+messages = [
+    # ... up to 50 messages ...
+]
+
+cf = ContextFlow()
+
+result = cf.optimize(
+    messages=messages,
+    agent_goal="Resolve customer shipping inquiry",  # this guides scoring
+    max_tokens=400  # your target for output
+)
+
+print("Tokens before:", result["analytics"]["tokens_before"])
+print("Tokens after:", result["analytics"]["tokens_after"])
+print("Reduction:", result["analytics"]["reduction_pct"], "%")
+print("Messages ready for LLM:", result["messages"])
+
+# Now use result.messages in your LLM query!
+# response = openai.ChatCompletion.create(model="gpt-4", messages=result.messages)
+```
+# License
+MIT License - see [LICENSE](https://github.com/athaapa/contextflow/blob/6b79c95d55c53c40bd8f500a6c921ea394c32666/LICENSE) file for details.

contextflow_ai-0.1.2.dist-info/RECORD

@@ -0,0 +1,16 @@
+contextflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+contextflow/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+contextflow/core/compactor.py,sha256=Sv_KflgD67Raujb2eIu9grv-irT4MCOYhNkhr-rWjD0,3911
+contextflow/core/optimizer.py,sha256=falRan-SLVZ-8a5i5V7R0ZSJEW16bmrNdMTV3RA9QbY,1195
+contextflow/core/scorer.py,sha256=CcSqasPa3-yPKJmIJUXwDCnuvaOh9COMRkcZ2FEXgas,2244
+contextflow/core/strategies.py,sha256=9d3j8sy7IEsmG7Hk45tiOY4axiSIjjAey9Q1I1jHMGM,4260
+contextflow/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+contextflow/utils/cost_calculator.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+contextflow/utils/embeddings.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+contextflow/utils/llm.py,sha256=ddVs9rwl-eJCmx9yUcR88KWpdv4-aESuIACO5l4PvA0,6138
+contextflow/utils/tokenizer.py,sha256=C3Bwo2w8aYMDts7tsBFlJCX8U8yhU5SKxWNOiiUqr8A,443
+contextflow_ai-0.1.2.dist-info/licenses/LICENSE,sha256=avSXTUv_48kpuVUFhhSvvklB-U4R5Pb9VcK-fuCP7zQ,1068
+contextflow_ai-0.1.2.dist-info/METADATA,sha256=X34X-nVW44k1-lLvHOS5U6L8QLkKCl_nysy621Wv1J4,2781
+contextflow_ai-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+contextflow_ai-0.1.2.dist-info/top_level.txt,sha256=zgLkleoXiZGCvGnFTWoxj7S_yLu8z_DC2qh_MstAkBs,12
+contextflow_ai-0.1.2.dist-info/RECORD,,

contextflow_ai-0.1.2.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

contextflow_ai-0.1.2.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Ayush Thapa
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

contextflow_ai-0.1.2.dist-info/top_level.txt

@@ -0,0 +1 @@
+contextflow