From 54d14dd319fc85cad1c3a3e32992c4be7b840378 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 00:42:08 +0200 Subject: [PATCH 01/14] Implement core functionality for COGNIX chat system with backend integration and conversation management --- cognix/__init__.py | 0 cognix/backend/__init__.py | 0 cognix/backend/base.py | 15 ++++++++++++++ cognix/backend/ollama_backend.py | 24 ++++++++++++++++++++++ cognix/conversation/__init__.py | 0 cognix/conversation/manager.py | 34 ++++++++++++++++++++++++++++++++ cognix/session.py | 24 ++++++++++++++++++++++ cognix/utils/__init__.py | 0 test_project_cognix/__init__.py | 0 test_project_cognix/main.py | 18 +++++++++++++++++ tests/__init__.py | 0 11 files changed, 115 insertions(+) create mode 100644 cognix/__init__.py create mode 100644 cognix/backend/__init__.py create mode 100644 cognix/backend/base.py create mode 100644 cognix/backend/ollama_backend.py create mode 100644 cognix/conversation/__init__.py create mode 100644 cognix/conversation/manager.py create mode 100644 cognix/session.py create mode 100644 cognix/utils/__init__.py create mode 100644 test_project_cognix/__init__.py create mode 100644 test_project_cognix/main.py create mode 100644 tests/__init__.py diff --git a/cognix/__init__.py b/cognix/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cognix/backend/__init__.py b/cognix/backend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cognix/backend/base.py b/cognix/backend/base.py new file mode 100644 index 0000000..a78782f --- /dev/null +++ b/cognix/backend/base.py @@ -0,0 +1,15 @@ +from abc import ABC, abstractmethod + +class BaseBackend(ABC): + @abstractmethod + def generate(self, prompt: str) -> str: + pass + +class OllamaBackend(BaseBackend): + def __init__(self, model): + self.model = model + # init connection stuff + + def generate(self, prompt: str) -> str: + # call ollama CLI or python lib here + return "ollama response" diff --git a/cognix/backend/ollama_backend.py b/cognix/backend/ollama_backend.py new file mode 100644 index 0000000..b7db0ae --- /dev/null +++ b/cognix/backend/ollama_backend.py @@ -0,0 +1,24 @@ +import subprocess +import json +from cognix.backend.base import BaseBackend + +class OllamaBackend(BaseBackend): + def __init__(self, model: str): + self.model = model + + def generate(self, prompt: str) -> str: + # Call ollama CLI to get response from the model + # Example: ollama generate --prompt "" + + try: + result = subprocess.run( + ["ollama", "run", self.model, prompt], + capture_output=True, + text=True, + check=True + ) + # Ollama CLI prints the response directly + return result.stdout.strip() + except subprocess.CalledProcessError as e: + # Handle errors cleanly + raise RuntimeError(f"Ollama generate failed: {e.stderr.strip()}") from e diff --git a/cognix/conversation/__init__.py b/cognix/conversation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cognix/conversation/manager.py b/cognix/conversation/manager.py new file mode 100644 index 0000000..62d790b --- /dev/null +++ b/cognix/conversation/manager.py @@ -0,0 +1,34 @@ +class ConversationManager: + def __init__(self, system_prompt=None, max_history=20): + self.system_prompt = system_prompt + self.max_history = max_history + self.history = [] + if system_prompt: + self.history.append({"role": "system", "content": system_prompt}) + + def add_user_message(self, content: str): + self.history.append({"role": "user", "content": content}) + self._trim_history() + + def add_assistant_message(self, content: str): + self.history.append({"role": "assistant", "content": content}) + self._trim_history() + + def _trim_history(self): + if len(self.history) > self.max_history + 1: + self.history = [self.history[0]] + self.history[-self.max_history:] + + def get_prompt(self) -> str: + prompt = "" + for msg in self.history: + prompt += f"{msg['role'].capitalize()}: {msg['content']}\n" + prompt += "Assistant: " + return prompt + + def reset(self): + self.history = [] + if self.system_prompt: + self.history.append({"role": "system", "content": self.system_prompt}) + + def get_history(self): + return self.history diff --git a/cognix/session.py b/cognix/session.py new file mode 100644 index 0000000..a036767 --- /dev/null +++ b/cognix/session.py @@ -0,0 +1,24 @@ +from cognix.backend.ollama_backend import OllamaBackend +from cognix.conversation.manager import ConversationManager + +class COGNIXSession: + def __init__(self, model="llama2", system_prompt=None): + self.backend = OllamaBackend(model=model) + self.conv_manager = ConversationManager(system_prompt=system_prompt) + + def chat(self, user_input: str) -> str: + self.conv_manager.add_user_message(user_input) + prompt = self.conv_manager.get_prompt() + response = self.backend.generate(prompt) + self.conv_manager.add_assistant_message(response) + return response + + def reset(self): + self.conv_manager.reset() + + def set_system_prompt(self, prompt: str): + self.conv_manager.system_prompt = prompt + self.reset() + + def get_history(self): + return self.conv_manager.get_history() diff --git a/cognix/utils/__init__.py b/cognix/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_project_cognix/__init__.py b/test_project_cognix/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_project_cognix/main.py b/test_project_cognix/main.py new file mode 100644 index 0000000..b5fb220 --- /dev/null +++ b/test_project_cognix/main.py @@ -0,0 +1,18 @@ +from cognix.session import COGNIXSession + +def main(): + session = COGNIXSession(model="llama3", system_prompt="You are a helpful assistant.") + + print("Welcome to COGNIX test chat. Type 'exit' to quit.\n") + + while True: + user_input = input("You: ") + if user_input.strip().lower() == "exit": + print("Exiting...") + break + + response = session.chat(user_input) + print(f"Assistant: {response}\n") + +if __name__ == "__main__": + main() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 From e7983bc43d5797a1ad8296c9090ee3409b8711dc Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 11:52:15 +0200 Subject: [PATCH 02/14] base conversational modules --- cognix/backend/base.py | 21 +++--- cognix/backend/ollama_backend.py | 77 +++++++++++++++----- cognix/conversation/buffer.py | 46 ++++++++++++ cognix/conversation/manager.py | 34 --------- cognix/conversation/summary.py | 82 ++++++++++++++++++++++ cognix/session.py | 33 +++++++-- cognix/utils/const/ENDPOINTS/llm.py | 2 + cognix/utils/const/PROMPTS/conversation.py | 9 +++ test_project_cognix/main.py | 26 +++++-- 9 files changed, 260 insertions(+), 70 deletions(-) create mode 100644 cognix/conversation/buffer.py create mode 100644 cognix/conversation/summary.py create mode 100644 cognix/utils/const/ENDPOINTS/llm.py create mode 100644 cognix/utils/const/PROMPTS/conversation.py diff --git a/cognix/backend/base.py b/cognix/backend/base.py index a78782f..1dc4626 100644 --- a/cognix/backend/base.py +++ b/cognix/backend/base.py @@ -2,14 +2,17 @@ class BaseBackend(ABC): @abstractmethod - def generate(self, prompt: str) -> str: + def generate(self, **kwargs) -> str: + """ + Generate a full response based on parameters. + Parameters depend on the specific backend implementation. + """ pass -class OllamaBackend(BaseBackend): - def __init__(self, model): - self.model = model - # init connection stuff - - def generate(self, prompt: str) -> str: - # call ollama CLI or python lib here - return "ollama response" + @abstractmethod + def generate_stream(self, **kwargs): + """ + Generate a streamed response based on parameters. + Yields chunks of the response. + """ + pass diff --git a/cognix/backend/ollama_backend.py b/cognix/backend/ollama_backend.py index b7db0ae..29f37f2 100644 --- a/cognix/backend/ollama_backend.py +++ b/cognix/backend/ollama_backend.py @@ -1,24 +1,67 @@ -import subprocess -import json +import requests from cognix.backend.base import BaseBackend +import json + +from cognix.utils.const.ENDPOINTS.llm import OLLAMA_API_URL_BASE, OLLAMA_API_URL_ENDPOINT class OllamaBackend(BaseBackend): - def __init__(self, model: str): + def __init__(self, model: str = None, api_url: str = OLLAMA_API_URL_BASE): self.model = model + self.api_url = api_url + + def generate(self, **kwargs) -> str: + url = self.api_url + OLLAMA_API_URL_ENDPOINT + + # Use model from kwargs or fallback to self.model + model = kwargs.pop("model", self.model) + if not model: + raise ValueError("Model name must be specified either at init or as a parameter.") + + # Compose payload with mandatory model and prompt + rest of kwargs + payload = { + "model": model, + **kwargs + } + + # Force stream=True to handle streaming internally but return concatenated string + payload["stream"] = True + try: + with requests.post(url, json=payload, stream=True) as response: + response.raise_for_status() + full_output = "" + for line in response.iter_lines(): + if line: + chunk = json.loads(line.decode("utf-8")) + full_output += chunk.get("response", "") + if chunk.get("done"): + break + return full_output.strip() + except requests.exceptions.RequestException as e: + raise RuntimeError(f"Ollama API request failed: {e}") + + def generate_stream(self, **kwargs): + url = f"{self.api_url}/api/generate" + + model = kwargs.pop("model", self.model) + if not model: + raise ValueError("Model name must be specified either at init or as a parameter.") - def generate(self, prompt: str) -> str: - # Call ollama CLI to get response from the model - # Example: ollama generate --prompt "" + payload = { + "model": model, + **kwargs + } + # Default to stream True if not specified + payload.setdefault("stream", True) try: - result = subprocess.run( - ["ollama", "run", self.model, prompt], - capture_output=True, - text=True, - check=True - ) - # Ollama CLI prints the response directly - return result.stdout.strip() - except subprocess.CalledProcessError as e: - # Handle errors cleanly - raise RuntimeError(f"Ollama generate failed: {e.stderr.strip()}") from e + with requests.post(url, json=payload, stream=payload["stream"]) as response: + response.raise_for_status() + for line in response.iter_lines(): + if line: + chunk = json.loads(line.decode("utf-8")) + if "response" in chunk: + yield chunk["response"] + if chunk.get("done"): + break + except requests.exceptions.RequestException as e: + raise RuntimeError(f"Ollama streaming API request failed: {e}") diff --git a/cognix/conversation/buffer.py b/cognix/conversation/buffer.py new file mode 100644 index 0000000..7455153 --- /dev/null +++ b/cognix/conversation/buffer.py @@ -0,0 +1,46 @@ +class BufferMemory: + def __init__(self, system_prompt=None, max_history=20): + self.system_prompt = system_prompt + self.max_history = max_history + self.history = [] + if system_prompt: + self.history.append({"role": "system", "content": system_prompt}) + + def add_user_message(self, content: str): + self.history.append({"role": "user", "content": content}) + self._trim_history() + + def add_assistant_message(self, content: str): + self.history.append({"role": "assistant", "content": content}) + self._trim_history() + + def _trim_history(self): + if len(self.history) > self.max_history + 1: + self.history = [self.history[0]] + self.history[-self.max_history:] + + def get_prompt(self) -> str: + conversation = "" + for msg in self.history: + if msg["role"] == "system": + continue # Don't include the system message in the chat history here + role = "Human" if msg["role"] == "user" else "AI" + conversation += f"{role}: {msg['content']}\n" + prompt = ( + f"{self.system_prompt.strip()}\n\n" if self.system_prompt else "" + ) + ( + "The following is a friendly conversation between a human and an AI. " + "The AI is talkative and provides lots of specific details from its context. " + "If the AI does not know the answer to a question, it truthfully says it does not know.\n\n" + f"Current conversation:\n{conversation}Human: " + ) + print(f"DEBUG: \n ############### \n {prompt} \n ############### \n") + return prompt + + + def reset(self): + self.history = [] + if self.system_prompt: + self.history.append({"role": "system", "content": self.system_prompt}) + + def get_history(self): + return self.history diff --git a/cognix/conversation/manager.py b/cognix/conversation/manager.py index 62d790b..e69de29 100644 --- a/cognix/conversation/manager.py +++ b/cognix/conversation/manager.py @@ -1,34 +0,0 @@ -class ConversationManager: - def __init__(self, system_prompt=None, max_history=20): - self.system_prompt = system_prompt - self.max_history = max_history - self.history = [] - if system_prompt: - self.history.append({"role": "system", "content": system_prompt}) - - def add_user_message(self, content: str): - self.history.append({"role": "user", "content": content}) - self._trim_history() - - def add_assistant_message(self, content: str): - self.history.append({"role": "assistant", "content": content}) - self._trim_history() - - def _trim_history(self): - if len(self.history) > self.max_history + 1: - self.history = [self.history[0]] + self.history[-self.max_history:] - - def get_prompt(self) -> str: - prompt = "" - for msg in self.history: - prompt += f"{msg['role'].capitalize()}: {msg['content']}\n" - prompt += "Assistant: " - return prompt - - def reset(self): - self.history = [] - if self.system_prompt: - self.history.append({"role": "system", "content": self.system_prompt}) - - def get_history(self): - return self.history diff --git a/cognix/conversation/summary.py b/cognix/conversation/summary.py new file mode 100644 index 0000000..7286112 --- /dev/null +++ b/cognix/conversation/summary.py @@ -0,0 +1,82 @@ +class SummaryMemory: + def __init__(self, summarizer, system_prompt=None, max_recent=4): + self.summarizer = summarizer + self.system_prompt = system_prompt + self.summary = "" # Keep summary separate from system_prompt + self.recent_history = [] + self.max_recent = max_recent + + def add_user_message(self, content: str): + self.recent_history.append({"role": "user", "content": content}) + self._prune_recent_entries() + + def add_assistant_message(self, content: str): + self.recent_history.append({"role": "assistant", "content": content}) + self._prune_recent_entries() + + def _prune_recent_entries(self): + if len(self.recent_history) > self.max_recent: + self._update_summary() + self.recent_history = self.recent_history[-self.max_recent:] + + + + def _format_new_lines(self) -> str: + lines = [] + for msg in self.recent_history: + if msg["role"] == "user": + lines.append(f"User:\n{msg['content']}") + else: + lines.append(f"{msg['content']}") + return "\n\n".join(lines) + + def _last_user_message(self) -> str: + for msg in reversed(self.recent_history): + if msg["role"] == "user": + return msg["content"] + return "" + def _last_assistant_message(self) -> str: + for msg in reversed(self.recent_history): + if msg["role"] == "assistant": + return msg["content"] + return "" + + def _update_summary(self): + new_lines = self._format_new_lines() + prompt = f"""You are an AI assistant tasked with summarizing a conversation between a User and an Assistant. Only summarize the Assistant's responses. + + Your goal is to produce a concise, comprehensive summary that captures all important details without losing any information. + The summary must be a single continuous block of text with no lists, comments, or formatting. + Do not include any explanations, quotes, or extra commentary. + + Current summary: + {self.summary.strip()} + + New dialogue: + {new_lines}""" + + self.summary = self.summarizer(prompt).strip() + + + + def get_prompt(self) -> str: + prompt_parts = [] + if self.system_prompt: + prompt_parts.append(f"[SYSTEM] {self.system_prompt.strip()}") + if self.summary: + prompt_parts.append(f"[ASSISTANT] {self.summary.strip()}") + if self.recent_history: + prompt_parts.append(f"[USER] {self._last_user_message()}") + prompt = "\n\n".join(prompt_parts) + + print(f"DEBUG:\n###############\n{prompt}\n###############\n") + return prompt + + + + def reset(self): + self.summary = "" + self.recent_history = [] + + def get_history(self): + return {"summary": self.summary} diff --git a/cognix/session.py b/cognix/session.py index a036767..7c64ff9 100644 --- a/cognix/session.py +++ b/cognix/session.py @@ -1,15 +1,28 @@ from cognix.backend.ollama_backend import OllamaBackend -from cognix.conversation.manager import ConversationManager +from cognix.conversation.buffer import BufferMemory +from cognix.conversation.summary import SummaryMemory + +from cognix.utils.const.PROMPTS.conversation import BASE_CONVESATION_PROMPT class COGNIXSession: - def __init__(self, model="llama2", system_prompt=None): + def __init__(self, model="llama2", system_prompt=BASE_CONVESATION_PROMPT, memory_type="summary"): self.backend = OllamaBackend(model=model) - self.conv_manager = ConversationManager(system_prompt=system_prompt) - def chat(self, user_input: str) -> str: + if memory_type == "summary": + self.conv_manager = SummaryMemory( + summarizer=self._summarizer, + system_prompt=system_prompt + ) + else: + self.conv_manager = BufferMemory(system_prompt=system_prompt) + + def _summarizer(self, prompt: str) -> str: + return self.backend.generate(prompt = prompt) + + def chat(self, user_input: str, **kwargs) -> str: self.conv_manager.add_user_message(user_input) prompt = self.conv_manager.get_prompt() - response = self.backend.generate(prompt) + response = self.backend.generate(prompt=prompt, **kwargs) self.conv_manager.add_assistant_message(response) return response @@ -22,3 +35,13 @@ def set_system_prompt(self, prompt: str): def get_history(self): return self.conv_manager.get_history() + + def chat_stream(self, user_input: str, **kwargs): + self.conv_manager.add_user_message(user_input) + prompt = self.conv_manager.get_prompt() + stream = self.backend.generate_stream(prompt=prompt, **kwargs) + full_response = "" + for chunk in stream: + full_response += chunk + yield chunk + self.conv_manager.add_assistant_message(full_response) diff --git a/cognix/utils/const/ENDPOINTS/llm.py b/cognix/utils/const/ENDPOINTS/llm.py new file mode 100644 index 0000000..e0e9803 --- /dev/null +++ b/cognix/utils/const/ENDPOINTS/llm.py @@ -0,0 +1,2 @@ +OLLAMA_API_URL_BASE = "http://localhost:11434" +OLLAMA_API_URL_ENDPOINT = "/api/generate" \ No newline at end of file diff --git a/cognix/utils/const/PROMPTS/conversation.py b/cognix/utils/const/PROMPTS/conversation.py new file mode 100644 index 0000000..e7a0947 --- /dev/null +++ b/cognix/utils/const/PROMPTS/conversation.py @@ -0,0 +1,9 @@ + + +BASE_CONVESATION_PROMPT = ( + "The following is a friendly, detailed, and informative conversation between a human and an AI assistant. " + "The AI only speaks for itself and never includes the human's part or pretends to be both sides. " + "Each AI reply is a continuation of the conversation without any prefixes like 'AI:' or 'Human:'. " + "The AI remains talkative, helpful, and provides specific, context-aware answers. " + "If the AI does not know something, it clearly states that it does not know.\n\n" +) diff --git a/test_project_cognix/main.py b/test_project_cognix/main.py index b5fb220..ffabbd2 100644 --- a/test_project_cognix/main.py +++ b/test_project_cognix/main.py @@ -1,18 +1,34 @@ from cognix.session import COGNIXSession def main(): - session = COGNIXSession(model="llama3", system_prompt="You are a helpful assistant.") + session = COGNIXSession() - print("Welcome to COGNIX test chat. Type 'exit' to quit.\n") + print("Welcome to COGNIX test chat. Type 'exit' to quit.") + stream_mode = True + + # Example extra parameters to test + extra_params = { + "raw": True, + } + + print() while True: - user_input = input("You: ") + user_input = input("#######################\n human input on test app : ") if user_input.strip().lower() == "exit": print("Exiting...") break - response = session.chat(user_input) - print(f"Assistant: {response}\n") + if stream_mode: + # Pass extra params to chat_stream + for chunk in session.chat_stream(user_input, **extra_params): + print(chunk, end="", flush=True) + print("\n#######################\n") # New line after stream + else: + # Pass extra params to chat + response = session.chat(user_input, **extra_params) + print(f"{response}") + print("\n") # New line after response if __name__ == "__main__": main() From b96814f62723bdbf845dbd575e2a16178e0a443a Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:25:53 +0200 Subject: [PATCH 03/14] Implement BufferMemory and SummaryMemory classes for conversation management --- .../{ => memory_management}/buffer.py | 0 .../memory_management/memory_strategy.py | 18 ++++ .../conversation/memory_management/summary.py | 70 ++++++++++++++++ cognix/conversation/summary.py | 82 ------------------- 4 files changed, 88 insertions(+), 82 deletions(-) rename cognix/conversation/{ => memory_management}/buffer.py (100%) create mode 100644 cognix/conversation/memory_management/memory_strategy.py create mode 100644 cognix/conversation/memory_management/summary.py delete mode 100644 cognix/conversation/summary.py diff --git a/cognix/conversation/buffer.py b/cognix/conversation/memory_management/buffer.py similarity index 100% rename from cognix/conversation/buffer.py rename to cognix/conversation/memory_management/buffer.py diff --git a/cognix/conversation/memory_management/memory_strategy.py b/cognix/conversation/memory_management/memory_strategy.py new file mode 100644 index 0000000..c1b5801 --- /dev/null +++ b/cognix/conversation/memory_management/memory_strategy.py @@ -0,0 +1,18 @@ +from abc import ABC, abstractmethod + +class MemoryStrategy(ABC): + @abstractmethod + def add_message(self, role: str, content: str): + pass + + @abstractmethod + def get_prompt(self) -> str: + pass + + @abstractmethod + def reset(self): + pass + + @abstractmethod + def get_state(self) -> dict: + pass diff --git a/cognix/conversation/memory_management/summary.py b/cognix/conversation/memory_management/summary.py new file mode 100644 index 0000000..5de2788 --- /dev/null +++ b/cognix/conversation/memory_management/summary.py @@ -0,0 +1,70 @@ +from cognix.conversation.memory_management.memory_strategy import MemoryStrategy +import threading + +class SummaryMemory(MemoryStrategy): + def __init__(self, summarizer, max_recent=4): + self.summarizer = summarizer + self.summary = "" + self.recent_history = [] + self.max_recent = max_recent + + def add_message(self, role: str, content: str): + self.recent_history.append({"role": role, "content": content}) + self._prune_recent_entries() + + def _prune_recent_entries(self): + if len(self.recent_history) > self.max_recent: + self._update_summary_async() + self.recent_history = [self._last_user_message()] + + def _last_user_message(self): + return next((msg for msg in reversed(self.recent_history) if msg["role"] == "user"), "") + + def _format_new_lines(self) -> str: + return "\n\n".join( + f"[USER] {m['content']}" if m["role"] == "user" else f"[CHATBOT] {m['content']}" + for m in self.recent_history + ) + + def _update_summary(self): + new_lines = self._format_new_lines() + prompt = f""" + Can you please summarize the following conversation and provide a brief overview of the main points discussed? + I want only the summary of the conversation, not the conversation itself. + {self.summary.strip()} + {new_lines}""" + self.summary = self.summarizer(prompt).strip() + + def _update_summary_async(self): + def summarization_task(): + new_lines = self._format_new_lines() + prompt = f""" + Can you please summarize the following conversation and provide a brief overview of the main points discussed? + I want only the summary of the conversation, not the conversation itself. + {self.summary.strip()} + {new_lines}""" + new_summary = self.summarizer(prompt).strip() + self.summary = new_summary # safely update shared state + + thread = threading.Thread(target=summarization_task, daemon=True) + thread.start() + + def get_prompt(self) -> str: + parts = [] + if self.summary: + parts.append(f"[SUMMARY] {self.summary.strip()}") + if self.recent_history: + parts.append(self._format_new_lines().strip()) + parts.append("\n[CHATBOT] ") + return "\n\n".join(parts) + + def reset(self): + self.summary = "" + self.recent_history = [] + + def get_state(self): + return {"summary": self.summary} + + +def summary_memory_factory(config): + return SummaryMemory(**config) \ No newline at end of file diff --git a/cognix/conversation/summary.py b/cognix/conversation/summary.py deleted file mode 100644 index 7286112..0000000 --- a/cognix/conversation/summary.py +++ /dev/null @@ -1,82 +0,0 @@ -class SummaryMemory: - def __init__(self, summarizer, system_prompt=None, max_recent=4): - self.summarizer = summarizer - self.system_prompt = system_prompt - self.summary = "" # Keep summary separate from system_prompt - self.recent_history = [] - self.max_recent = max_recent - - def add_user_message(self, content: str): - self.recent_history.append({"role": "user", "content": content}) - self._prune_recent_entries() - - def add_assistant_message(self, content: str): - self.recent_history.append({"role": "assistant", "content": content}) - self._prune_recent_entries() - - def _prune_recent_entries(self): - if len(self.recent_history) > self.max_recent: - self._update_summary() - self.recent_history = self.recent_history[-self.max_recent:] - - - - def _format_new_lines(self) -> str: - lines = [] - for msg in self.recent_history: - if msg["role"] == "user": - lines.append(f"User:\n{msg['content']}") - else: - lines.append(f"{msg['content']}") - return "\n\n".join(lines) - - def _last_user_message(self) -> str: - for msg in reversed(self.recent_history): - if msg["role"] == "user": - return msg["content"] - return "" - def _last_assistant_message(self) -> str: - for msg in reversed(self.recent_history): - if msg["role"] == "assistant": - return msg["content"] - return "" - - def _update_summary(self): - new_lines = self._format_new_lines() - prompt = f"""You are an AI assistant tasked with summarizing a conversation between a User and an Assistant. Only summarize the Assistant's responses. - - Your goal is to produce a concise, comprehensive summary that captures all important details without losing any information. - The summary must be a single continuous block of text with no lists, comments, or formatting. - Do not include any explanations, quotes, or extra commentary. - - Current summary: - {self.summary.strip()} - - New dialogue: - {new_lines}""" - - self.summary = self.summarizer(prompt).strip() - - - - def get_prompt(self) -> str: - prompt_parts = [] - if self.system_prompt: - prompt_parts.append(f"[SYSTEM] {self.system_prompt.strip()}") - if self.summary: - prompt_parts.append(f"[ASSISTANT] {self.summary.strip()}") - if self.recent_history: - prompt_parts.append(f"[USER] {self._last_user_message()}") - prompt = "\n\n".join(prompt_parts) - - print(f"DEBUG:\n###############\n{prompt}\n###############\n") - return prompt - - - - def reset(self): - self.summary = "" - self.recent_history = [] - - def get_history(self): - return {"summary": self.summary} From 4e63d7f87bf7be0b050a9ffe6fd766c1613d6289 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:25:59 +0200 Subject: [PATCH 04/14] Add __init__.py file to initialize the agents module --- cognix/agents/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 cognix/agents/__init__.py diff --git a/cognix/agents/__init__.py b/cognix/agents/__init__.py new file mode 100644 index 0000000..e69de29 From aa1e732ce9329166185c764da4ae124ae287e980 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:26:05 +0200 Subject: [PATCH 05/14] Add __init__.py file to initialize the summarizer module --- cognix/agents/summarizer/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 cognix/agents/summarizer/__init__.py diff --git a/cognix/agents/summarizer/__init__.py b/cognix/agents/summarizer/__init__.py new file mode 100644 index 0000000..e69de29 From d5046f2b44eacf7c975091b9c30182cbc2e6ebb9 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:26:07 +0200 Subject: [PATCH 06/14] Add __init__.py file to initialize the memory_management module --- cognix/conversation/memory_management/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 cognix/conversation/memory_management/__init__.py diff --git a/cognix/conversation/memory_management/__init__.py b/cognix/conversation/memory_management/__init__.py new file mode 100644 index 0000000..e69de29 From 3912732ba96d3c91c457e8c0d870686a971c6cfa Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:26:35 +0200 Subject: [PATCH 07/14] Add configuration constants for summary and buffer management --- cognix/utils/const/config.py | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 cognix/utils/const/config.py diff --git a/cognix/utils/const/config.py b/cognix/utils/const/config.py new file mode 100644 index 0000000..703c237 --- /dev/null +++ b/cognix/utils/const/config.py @@ -0,0 +1,9 @@ + + +SUMMARY_CONFIG = { + "max_recent": 5, +} + +BUFFER_CONFIG = { + "max_history": 20, +} \ No newline at end of file From 12f3041f8573f4ccf412c718c12efc9a1f53d3ea Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:26:38 +0200 Subject: [PATCH 08/14] Add MemoryRegistry class for managing memory types and factories --- cognix/utils/const/registry.py | 38 ++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 cognix/utils/const/registry.py diff --git a/cognix/utils/const/registry.py b/cognix/utils/const/registry.py new file mode 100644 index 0000000..9012463 --- /dev/null +++ b/cognix/utils/const/registry.py @@ -0,0 +1,38 @@ +# constants/registry.py + +from cognix.conversation.memory_management.summary import summary_memory_factory +from cognix.conversation.memory_management.buffer import buffer_memory_factory + + + +class MemoryRegistry: + _registry = {} + + @classmethod + def register(cls, name): + def wrapper(factory_fn): + cls._registry[name] = factory_fn + return factory_fn + return wrapper + + @classmethod + def get(cls, name, config): + if name not in cls._registry: + raise ValueError(f"Memory type '{name}' not found.") + return cls._registry[name](config) + + @classmethod + def list(cls): + return list(cls._registry.keys()) + + + +MEMORY_REGISTRY = { + "summary": lambda config: SummaryMemory( + summarizer=config["summarizer"], + config=config + ), + "buffer": lambda config: BufferMemory(config=config), +} +MemoryRegistry.register("summary")(summary_memory_factory) +MemoryRegistry.register("buffer")(buffer_memory_factory) From f1f80d452b0247b46472e8dc8cb95dbae1364dd1 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:26:58 +0200 Subject: [PATCH 09/14] Refactor BufferMemory class to inherit from MemoryStrategy and streamline message handling --- .../conversation/memory_management/buffer.py | 61 ++++++++++++------- 1 file changed, 38 insertions(+), 23 deletions(-) diff --git a/cognix/conversation/memory_management/buffer.py b/cognix/conversation/memory_management/buffer.py index 7455153..c70a941 100644 --- a/cognix/conversation/memory_management/buffer.py +++ b/cognix/conversation/memory_management/buffer.py @@ -1,46 +1,61 @@ -class BufferMemory: - def __init__(self, system_prompt=None, max_history=20): +from cognix.conversation.memory_management.memory_strategy import MemoryStrategy + +from cognix.utils.const.config import BUFFER_CONFIG + +class BufferMemory(MemoryStrategy): + def __init__(self, system_prompt: str = None, max_history: int = 20): self.system_prompt = system_prompt self.max_history = max_history self.history = [] if system_prompt: self.history.append({"role": "system", "content": system_prompt}) - def add_user_message(self, content: str): - self.history.append({"role": "user", "content": content}) + def add_message(self, role: str, content: str): + if role not in {"user", "assistant"}: + raise ValueError("Role must be 'user' or 'assistant'") + self.history.append({"role": role, "content": content}) self._trim_history() - def add_assistant_message(self, content: str): - self.history.append({"role": "assistant", "content": content}) - self._trim_history() + def add_user_message(self, content: str): + self.add_message("user", content) - def _trim_history(self): - if len(self.history) > self.max_history + 1: - self.history = [self.history[0]] + self.history[-self.max_history:] + def add_assistant_message(self, content: str): + self.add_message("assistant", content) def get_prompt(self) -> str: conversation = "" for msg in self.history: if msg["role"] == "system": - continue # Don't include the system message in the chat history here - role = "Human" if msg["role"] == "user" else "AI" - conversation += f"{role}: {msg['content']}\n" - prompt = ( - f"{self.system_prompt.strip()}\n\n" if self.system_prompt else "" - ) + ( - "The following is a friendly conversation between a human and an AI. " - "The AI is talkative and provides lots of specific details from its context. " - "If the AI does not know the answer to a question, it truthfully says it does not know.\n\n" - f"Current conversation:\n{conversation}Human: " - ) - print(f"DEBUG: \n ############### \n {prompt} \n ############### \n") + continue + role = "[USER] " if msg["role"] == "user" else "[CHATBOT] " + conversation += f"{role}{msg['content']}\n" + prompt = conversation.strip() + "\n[CHATBOT] " return prompt - def reset(self): self.history = [] if self.system_prompt: self.history.append({"role": "system", "content": self.system_prompt}) + def get_state(self) -> dict: + return {"history": self.history} + def get_history(self): return self.history + + def _trim_history(self): + system_msg = self.history[0] if self.history and self.history[0]["role"] == "system" else None + messages = self.history[1:] if system_msg else self.history + trimmed = messages[-self.max_history:] + self.history = [system_msg] + trimmed if system_msg else trimmed + + +def buffer_memory_factory(config=BUFFER_CONFIG) -> BufferMemory: + """ + Factory function to create a BufferMemory instance. + :param config: Configuration dictionary for BufferMemory. + :return: An instance of BufferMemory. + """ + if config is None: + config = BUFFER_CONFIG + return BufferMemory(**config) From edf469ca0a25eff1474552002dcd3243b1cda26a Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:27:03 +0200 Subject: [PATCH 10/14] Refactor SummaryMemory initialization and update summary prompt for clarity --- .../conversation/memory_management/summary.py | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/cognix/conversation/memory_management/summary.py b/cognix/conversation/memory_management/summary.py index 5de2788..757a463 100644 --- a/cognix/conversation/memory_management/summary.py +++ b/cognix/conversation/memory_management/summary.py @@ -1,9 +1,15 @@ +from cognix.utils.const.config import SUMMARY_CONFIG + + from cognix.conversation.memory_management.memory_strategy import MemoryStrategy import threading class SummaryMemory(MemoryStrategy): - def __init__(self, summarizer, max_recent=4): - self.summarizer = summarizer + def __init__(self, summarizer= None, max_recent: int = 5): + if summarizer is None: + self.summarizer = self._summarizer + else: + self.summarizer = summarizer self.summary = "" self.recent_history = [] self.max_recent = max_recent @@ -40,7 +46,7 @@ def summarization_task(): new_lines = self._format_new_lines() prompt = f""" Can you please summarize the following conversation and provide a brief overview of the main points discussed? - I want only the summary of the conversation, not the conversation itself. + I want only the summary of the conversation, not the conversation itself. Do not include the [USER] and [CHATBOT] tags. {self.summary.strip()} {new_lines}""" new_summary = self.summarizer(prompt).strip() @@ -48,6 +54,13 @@ def summarization_task(): thread = threading.Thread(target=summarization_task, daemon=True) thread.start() + + def _summarizer(self, prompt: str) -> str: + # TODO FIX THIS SHIT + # temporary placeholder for the summarizer agent that will be used + from cognix.backend.ollama_backend import OllamaBackend + backend = OllamaBackend(model="llama2") + return backend.generate(prompt=prompt) def get_prompt(self) -> str: parts = [] @@ -66,5 +79,10 @@ def get_state(self): return {"summary": self.summary} -def summary_memory_factory(config): +def summary_memory_factory(config=SUMMARY_CONFIG) -> SummaryMemory: + """ + Factory function to create a SummaryMemory instance. + """ + if config is None: + config = SUMMARY_CONFIG return SummaryMemory(**config) \ No newline at end of file From e0945862148076ac28b4d5d327b00cbced7f3f75 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:27:09 +0200 Subject: [PATCH 11/14] Implement ConversationManager class for managing user and assistant messages --- cognix/conversation/manager.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/cognix/conversation/manager.py b/cognix/conversation/manager.py index e69de29..a3fb871 100644 --- a/cognix/conversation/manager.py +++ b/cognix/conversation/manager.py @@ -0,0 +1,21 @@ +from cognix.conversation.memory_management.memory_strategy import MemoryStrategy + +class ConversationManager: + def __init__(self, memory: MemoryStrategy, system_prompt=None): + self.memory = memory + self.system_prompt = system_prompt + + def add_user_message(self, content: str): + self.memory.add_message("user", content) + + def add_assistant_message(self, content: str): + self.memory.add_message("assistant", content) + + def get_prompt(self) -> str: + return self.memory.get_prompt() + + def reset_conversation(self): + self.memory.reset() + + def export_state(self): + return self.memory.get_state() From 6baebb8de5503f0475662dbbd44692d5b160375a Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:27:12 +0200 Subject: [PATCH 12/14] Refactor COGNIXSession to use MemoryRegistry for dynamic memory management and streamline initialization --- cognix/session.py | 42 ++++++++++++++++++++---------------------- 1 file changed, 20 insertions(+), 22 deletions(-) diff --git a/cognix/session.py b/cognix/session.py index 7c64ff9..59fd6f3 100644 --- a/cognix/session.py +++ b/cognix/session.py @@ -1,23 +1,21 @@ from cognix.backend.ollama_backend import OllamaBackend -from cognix.conversation.buffer import BufferMemory -from cognix.conversation.summary import SummaryMemory + +from cognix.conversation.memory_management.buffer import BufferMemory +from cognix.conversation.memory_management.summary import SummaryMemory +from cognix.conversation.manager import ConversationManager from cognix.utils.const.PROMPTS.conversation import BASE_CONVESATION_PROMPT +from cognix.utils.const.registry import MemoryRegistry class COGNIXSession: - def __init__(self, model="llama2", system_prompt=BASE_CONVESATION_PROMPT, memory_type="summary"): + def __init__(self, model="llama2", system_prompt=BASE_CONVESATION_PROMPT, memory_type="summary",memory_config=None): self.backend = OllamaBackend(model=model) - if memory_type == "summary": - self.conv_manager = SummaryMemory( - summarizer=self._summarizer, - system_prompt=system_prompt - ) - else: - self.conv_manager = BufferMemory(system_prompt=system_prompt) + # Use the registry to get the memory instance dynamically + memory = MemoryRegistry.get(memory_type, memory_config) + + self.conv_manager = ConversationManager(memory=memory, system_prompt=system_prompt) - def _summarizer(self, prompt: str) -> str: - return self.backend.generate(prompt = prompt) def chat(self, user_input: str, **kwargs) -> str: self.conv_manager.add_user_message(user_input) @@ -25,6 +23,16 @@ def chat(self, user_input: str, **kwargs) -> str: response = self.backend.generate(prompt=prompt, **kwargs) self.conv_manager.add_assistant_message(response) return response + + def chat_stream(self, user_input: str, **kwargs): + self.conv_manager.add_user_message(user_input) + prompt = self.conv_manager.get_prompt() + stream = self.backend.generate_stream(prompt=prompt, **kwargs) + full_response = "" + for chunk in stream: + full_response += chunk + yield chunk + self.conv_manager.add_assistant_message(full_response) def reset(self): self.conv_manager.reset() @@ -35,13 +43,3 @@ def set_system_prompt(self, prompt: str): def get_history(self): return self.conv_manager.get_history() - - def chat_stream(self, user_input: str, **kwargs): - self.conv_manager.add_user_message(user_input) - prompt = self.conv_manager.get_prompt() - stream = self.backend.generate_stream(prompt=prompt, **kwargs) - full_response = "" - for chunk in stream: - full_response += chunk - yield chunk - self.conv_manager.add_assistant_message(full_response) From 951e3964bc811fe3d9cad69937bf2ca064e0bc42 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:27:16 +0200 Subject: [PATCH 13/14] Add new agent.py file for summarizer functionality --- cognix/agents/summarizer/agent.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 cognix/agents/summarizer/agent.py diff --git a/cognix/agents/summarizer/agent.py b/cognix/agents/summarizer/agent.py new file mode 100644 index 0000000..e69de29 From c0b621074157bc701c913206aa9a0727597fe341 Mon Sep 17 00:00:00 2001 From: comus3 Date: Tue, 20 May 2025 17:27:20 +0200 Subject: [PATCH 14/14] Refactor main.py to enhance user input display with color and streamline input handling --- test_project_cognix/main.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test_project_cognix/main.py b/test_project_cognix/main.py index ffabbd2..fec1be1 100644 --- a/test_project_cognix/main.py +++ b/test_project_cognix/main.py @@ -1,5 +1,5 @@ from cognix.session import COGNIXSession - +from colorama import Fore, Style def main(): session = COGNIXSession() @@ -11,10 +11,11 @@ def main(): "raw": True, } - print() - + # input in green method + def input_in_green(prompt): + return input(f"{Fore.GREEN}{prompt}{Style.RESET_ALL}") while True: - user_input = input("#######################\n human input on test app : ") + user_input = input_in_green("\nhuman input on test app : ") if user_input.strip().lower() == "exit": print("Exiting...") break @@ -23,7 +24,6 @@ def main(): # Pass extra params to chat_stream for chunk in session.chat_stream(user_input, **extra_params): print(chunk, end="", flush=True) - print("\n#######################\n") # New line after stream else: # Pass extra params to chat response = session.chat(user_input, **extra_params)