diff --git a/cognix/__init__.py b/cognix/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cognix/agents/__init__.py b/cognix/agents/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cognix/agents/summarizer/__init__.py b/cognix/agents/summarizer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cognix/agents/summarizer/agent.py b/cognix/agents/summarizer/agent.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..1dc4626 --- /dev/null +++ b/cognix/backend/base.py @@ -0,0 +1,18 @@ +from abc import ABC, abstractmethod + +class BaseBackend(ABC): + @abstractmethod + def generate(self, **kwargs) -> str: + """ + Generate a full response based on parameters. + Parameters depend on the specific backend implementation. + """ + pass + + @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 new file mode 100644 index 0000000..29f37f2 --- /dev/null +++ b/cognix/backend/ollama_backend.py @@ -0,0 +1,67 @@ +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 = 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.") + + payload = { + "model": model, + **kwargs + } + + # Default to stream True if not specified + payload.setdefault("stream", True) + try: + 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/__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..a3fb871 --- /dev/null +++ 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() diff --git a/cognix/conversation/memory_management/__init__.py b/cognix/conversation/memory_management/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cognix/conversation/memory_management/buffer.py b/cognix/conversation/memory_management/buffer.py new file mode 100644 index 0000000..c70a941 --- /dev/null +++ b/cognix/conversation/memory_management/buffer.py @@ -0,0 +1,61 @@ +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_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_user_message(self, content: str): + self.add_message("user", content) + + 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 + 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) 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..757a463 --- /dev/null +++ b/cognix/conversation/memory_management/summary.py @@ -0,0 +1,88 @@ +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= 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 + + 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. Do not include the [USER] and [CHATBOT] tags. + {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 _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 = [] + 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=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 diff --git a/cognix/session.py b/cognix/session.py new file mode 100644 index 0000000..59fd6f3 --- /dev/null +++ b/cognix/session.py @@ -0,0 +1,45 @@ +from cognix.backend.ollama_backend import OllamaBackend + +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",memory_config=None): + self.backend = OllamaBackend(model=model) + + # 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 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=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() + + 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/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/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 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) 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..fec1be1 --- /dev/null +++ b/test_project_cognix/main.py @@ -0,0 +1,34 @@ +from cognix.session import COGNIXSession +from colorama import Fore, Style +def main(): + session = COGNIXSession() + + print("Welcome to COGNIX test chat. Type 'exit' to quit.") + stream_mode = True + + # Example extra parameters to test + extra_params = { + "raw": True, + } + + # input in green method + def input_in_green(prompt): + return input(f"{Fore.GREEN}{prompt}{Style.RESET_ALL}") + while True: + user_input = input_in_green("\nhuman input on test app : ") + if user_input.strip().lower() == "exit": + print("Exiting...") + break + + if stream_mode: + # Pass extra params to chat_stream + for chunk in session.chat_stream(user_input, **extra_params): + print(chunk, end="", flush=True) + 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() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29