%pip install --quiet crewai langchain-couchbase langchain-openai python-dotenv

Note: you may need to restart the kernel to use updated packages.

from typing import Any, Dict, List, Optional
import os
import logging
from datetime import timedelta
from dotenv import load_dotenv
from crewai.memory.storage.rag_storage import RAGStorage
from crewai.memory.short_term.short_term_memory import ShortTermMemory
from crewai import Agent, Crew, Task, Process
from couchbase.cluster import Cluster
from couchbase.options import ClusterOptions
from couchbase.auth import PasswordAuthenticator
from couchbase.diagnostics import PingState, ServiceType
from langchain_couchbase.vectorstores import CouchbaseVectorStore
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
import time
import json
import uuid

# Configure logging (disabled)
logging.basicConfig(level=logging.CRITICAL)
logger = logging.getLogger(__name__)

load_dotenv()

# Verify environment variables
required_vars = ['OPENAI_API_KEY', 'CB_HOST', 'CB_USERNAME', 'CB_PASSWORD']
for var in required_vars:
    if not os.getenv(var):
        raise ValueError(f"{var} environment variable is required")

class CouchbaseStorage(RAGStorage):
    """
    Extends RAGStorage to handle embeddings for memory entries using Couchbase.
    """

    def __init__(self, type: str, allow_reset: bool = True, embedder_config: Optional[Dict[str, Any]] = None, crew: Optional[Any] = None):
        """Initialize CouchbaseStorage with configuration."""
        super().__init__(type, allow_reset, embedder_config, crew)
        self._initialize_app()

    def search(
        self,
        query: str,
        limit: int = 3,
        filter: Optional[dict] = None,
        score_threshold: float = 0,
    ) -> List[Dict[str, Any]]:
        """
        Search memory entries using vector similarity.
        """
        try:
            # Add type filter
            search_filter = {"memory_type": self.type}
            if filter:
                search_filter.update(filter)

            # Execute search
            results = self.vector_store.similarity_search_with_score(
                query,
                k=limit,
                filter=search_filter
            )
            
            # Format results and deduplicate by content
            seen_contents = set()
            formatted_results = []
            
            for i, (doc, score) in enumerate(results):
                if score >= score_threshold:
                    content = doc.page_content
                    if content not in seen_contents:
                        seen_contents.add(content)
                        formatted_results.append({
                            "id": doc.metadata.get("memory_id", str(i)),
                            "metadata": doc.metadata,
                            "context": content,
                            "score": float(score)
                        })
            
            logger.info(f"Found {len(formatted_results)} unique results for query: {query}")
            return formatted_results

        except Exception as e:
            logger.error(f"Search failed: {str(e)}")
            return []

    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
        """
        Save a memory entry with metadata.
        """
        try:
            # Generate unique ID
            memory_id = str(uuid.uuid4())
            timestamp = int(time.time() * 1000)
            
            # Prepare metadata (create a copy to avoid modifying references)
            if not metadata:
                metadata = {}
            else:
                metadata = metadata.copy()  # Create a copy to avoid modifying references
                
            # Process agent-specific information if present
            agent_name = metadata.get('agent', 'unknown')
                
            # Clean up value if it has the typical LLM response format
            value_str = str(value)
            if "Final Answer:" in value_str:
                # Extract just the actual content - everything after "Final Answer:"
                parts = value_str.split("Final Answer:", 1)
                if len(parts) > 1:
                    value = parts[1].strip()
                    logger.info(f"Cleaned up response format for agent: {agent_name}")
            elif value_str.startswith("Thought:"):
                # Handle thought/final answer format
                if "Final Answer:" in value_str:
                    parts = value_str.split("Final Answer:", 1)
                    if len(parts) > 1:
                        value = parts[1].strip()
                        logger.info(f"Cleaned up thought process format for agent: {agent_name}")
            
            # Update metadata
            metadata.update({
                "memory_id": memory_id,
                "memory_type": self.type,
                "timestamp": timestamp,
                "source": "crewai"
            })

            # Log memory information for debugging
            value_preview = str(value)[:100] + "..." if len(str(value)) > 100 else str(value)
            metadata_preview = {k: v for k, v in metadata.items() if k != "embedding"}
            logger.info(f"Saving memory for Agent: {agent_name}")
            logger.info(f"Memory value preview: {value_preview}")
            logger.info(f"Memory metadata: {metadata_preview}")
            
            # Convert value to string if needed
            if isinstance(value, (dict, list)):
                value = json.dumps(value)
            elif not isinstance(value, str):
                value = str(value)

            # Save to vector store
            self.vector_store.add_texts(
                texts=[value],
                metadatas=[metadata],
                ids=[memory_id]
            )
            logger.info(f"Saved memory {memory_id}: {value[:100]}...")

        except Exception as e:
            logger.error(f"Save failed: {str(e)}")
            raise

    def reset(self) -> None:
        """Reset the memory storage if allowed."""
        if not self.allow_reset:
            return

        try:
            # Delete documents of this memory type
            self.cluster.query(
                f"DELETE FROM `{self.bucket_name}`.`{self.scope_name}`.`{self.collection_name}` WHERE memory_type = $type",
                type=self.type
            ).execute()
            logger.info(f"Reset memory type: {self.type}")
        except Exception as e:
            logger.error(f"Reset failed: {str(e)}")
            raise

    def _initialize_app(self):
        """Initialize Couchbase connection and vector store."""
        try:
            # Initialize embeddings
            if self.embedder_config and self.embedder_config.get("provider") == "openai":
                self.embeddings = OpenAIEmbeddings(
                    openai_api_key=os.getenv('OPENAI_API_KEY'),
                    model=self.embedder_config.get("config", {}).get("model", "text-embedding-3-small")
                )
            else:
                self.embeddings = OpenAIEmbeddings(
                    openai_api_key=os.getenv('OPENAI_API_KEY'),
                    model="text-embedding-3-small"
                )

            # Connect to Couchbase
            auth = PasswordAuthenticator(
                os.getenv('CB_USERNAME', ''),
                os.getenv('CB_PASSWORD', '')
            )
            options = ClusterOptions(auth)
            
            # Initialize cluster connection
            self.cluster = Cluster(os.getenv('CB_HOST', ''), options)
            self.cluster.wait_until_ready(timedelta(seconds=5))

            # Check search service
            ping_result = self.cluster.ping()
            search_available = False
            for service_type, endpoints in ping_result.endpoints.items():
                if service_type == ServiceType.Search:
                    for endpoint in endpoints:
                        if endpoint.state == PingState.OK:
                            search_available = True
                            logger.info(f"Search service is responding at: {endpoint.remote}")
                            break
                    break
            if not search_available:
                raise RuntimeError("Search/FTS service not found or not responding")
            
            # Set up storage configuration
            self.bucket_name = os.getenv('CB_BUCKET_NAME', 'vector-search-testing')
            self.scope_name = os.getenv('SCOPE_NAME', 'shared')
            self.collection_name = os.getenv('COLLECTION_NAME', 'crew')
            self.index_name = os.getenv('INDEX_NAME', 'vector_search_crew')

            # Initialize vector store
            self.vector_store = CouchbaseVectorStore(
                cluster=self.cluster,
                bucket_name=self.bucket_name,
                scope_name=self.scope_name,
                collection_name=self.collection_name,
                embedding=self.embeddings,
                index_name=self.index_name
            )
            logger.info(f"Initialized CouchbaseStorage for type: {self.type}")

        except Exception as e:
            logger.error(f"Initialization failed: {str(e)}")
            raise

# Initialize storage
storage = CouchbaseStorage(
    type="short_term",
    embedder_config={
        "provider": "openai",
        "config": {"model": "text-embedding-3-small"}
    }
)

# Reset storage
storage.reset()

# Test storage
test_memory = "Pep Guardiola praised Manchester City's current form, saying 'The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased.'"
test_metadata = {"category": "sports", "test": "initial_memory"}
storage.save(test_memory, test_metadata)

# Test search
results = storage.search("What did Guardiola say about Manchester City?", limit=1)
for result in results:
    print(f"Found: {result['context']}\nScore: {result['score']}\nMetadata: {result['metadata']}")

Found: Pep Guardiola praised Manchester City's current form, saying 'The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased.'
Score: 0.6598483920097351
Metadata: {}

# Initialize ShortTermMemory with our storage
memory = ShortTermMemory(storage=storage)

# Initialize language model
llm = ChatOpenAI(
    model="gpt-4o",
    temperature=0.7
)

# Create agents with memory
sports_analyst = Agent(
    role='Sports Analyst',
    goal='Analyze Manchester City performance',
    backstory='Expert at analyzing football teams and providing insights on their performance',
    llm=llm,
    memory=True,
    memory_storage=memory
)

journalist = Agent(
    role='Sports Journalist',
    goal='Create engaging football articles',
    backstory='Experienced sports journalist who specializes in Premier League coverage',
    llm=llm,
    memory=True,
    memory_storage=memory
)

# Create tasks
analysis_task = Task(
    description='Analyze Manchester City\'s recent performance based on Pep Guardiola\'s comments: "The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased."',
    agent=sports_analyst,
    expected_output="A comprehensive analysis of Manchester City's current form based on Guardiola's comments."
)

writing_task = Task(
    description='Write a sports article about Manchester City\'s form using the analysis and Guardiola\'s comments.',
    agent=journalist,
    context=[analysis_task],
    expected_output="An engaging sports article about Manchester City's current form and Guardiola's perspective."
)

# Create crew with memory
crew = Crew(
    agents=[sports_analyst, journalist],
    tasks=[analysis_task, writing_task],
    process=Process.sequential,
    memory=True,
    short_term_memory=memory,  # Explicitly pass our memory implementation
    verbose=True
)

# Run the crew
result = crew.kickoff()

print("\nCrew Result:")
print("-" * 80)
print(result)
print("-" * 80)

# Agent: Sports Analyst
## Task: Analyze Manchester City's recent performance based on Pep Guardiola's comments: "The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased."


# Agent: Sports Analyst
## Final Answer: 
Analyzing Manchester City's recent performance based on Pep Guardiola's comments, it is clear that the team is operating at a high level. Guardiola's statement, "The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased," gives insight into both the current form and the overall environment within the squad.

Historically, Manchester City under Guardiola has been known for their possession-based style, high pressing, and tactical flexibility. Comparing their current form to past performances, City appears to be continuing their trend of dominance, especially in terms of controlling matches and creating goal-scoring opportunities. Recent match statistics highlight their effectiveness; for example, in their last few games, City has consistently maintained possession percentages above 60%, showcasing their ability to control the tempo of the game.

Key players such as Kevin De Bruyne, Erling Haaland, and Phil Foden have been pivotal in their recent successes. De Bruyne's vision and passing ability have been instrumental in breaking down defenses, while Haaland's goal-scoring prowess has added a new dimension to City's attacking threat. Foden's versatility and creativity have also been crucial, providing Guardiola with multiple tactical options.

Despite the positive comments from Guardiola, there are still areas where City can improve. Defensive lapses have occasionally been a concern, as seen in matches where they conceded goals due to individual errors or lapses in concentration. Addressing these issues will be critical for maintaining their form throughout the season.

Guardiola's satisfaction with the team's training and performance indicates a harmonious and focused squad, which is essential for sustained success. His comments also reflect the strategic innovations he has implemented, such as adapting formations and player roles to maximize the team's strengths and exploit opponents' weaknesses.

In conclusion, Manchester City's current form is a testament to Guardiola's tactical expertise and the players' execution on the pitch. While there are minor areas for improvement, the overall outlook is positive, and they are well-positioned to compete for titles this season. The combination of Guardiola's strategic leadership and the team's depth and quality suggests that City will continue to be a formidable force in football.


# Agent: Sports Journalist
## Task: Write a sports article about Manchester City's form using the analysis and Guardiola's comments.


# Agent: Sports Journalist
## Final Answer: 
**Manchester City: A Blueprint of Excellence in Current Form Under Guardiola**

In the ever-dynamic world of Premier League football, Manchester City’s form this season stands as a testament to their consistent excellence and tactical acuity under the stewardship of Pep Guardiola. As the footballing world watches, City’s recent performances have not only reasserted their domestic dominance but have also sent a signal of their intent on the European stage.

Pep Guardiola, the cerebral architect behind Manchester City’s success, recently voiced his satisfaction with the team’s current standing. "The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased," he remarked, encapsulating the harmony and precision that have become synonymous with his tenure at the club.

Historically, Guardiola's sides have been renowned for their possession-based play and tactical flexibility. This season, City continues to exemplify these principles, maintaining an impressive average possession of over 60% in recent matches. This ability to control the game’s tempo has been pivotal in their recent string of victories and is indicative of their strategic dominance on the pitch.

Key to this success are the stellar performances from City’s roster of world-class talent. Kevin De Bruyne, with his unparalleled vision and precision passing, has been instrumental in orchestrating play and dismantling defenses. Erling Haaland, a new addition to the squad, has bolstered City’s attacking threat with his prolific goal-scoring ability, adding a new dimension to their already potent offensive setup. Meanwhile, Phil Foden’s versatility and creativity on the field provide Guardiola with a plethora of tactical options, making him an invaluable asset.

Despite the glowing appraisal, Guardiola remains pragmatic, acknowledging areas for refinement. Defensive lapses, often due to individual errors or brief lapses in concentration, have been points of concern. Addressing these issues is crucial for maintaining their momentum and ensuring that they remain competitive across all fronts.

Guardiola’s comments reflect not only a satisfied manager but also a strategic thinker constantly innovating to maximize the team’s strengths. His ability to adapt formations and player roles is a hallmark of his managerial prowess, ensuring that City not only plays to their strengths but also exploits the weaknesses of their opponents.

Looking ahead, Manchester City faces a challenging fixture schedule, with both domestic and European competitions testing their resolve. However, with Guardiola at the helm and the squad’s depth and quality, they are well-positioned to surmount these challenges. The combination of strategic leadership and player execution suggests that City will continue to be a formidable force in football.

In conclusion, Manchester City’s current form is a masterclass in footballing excellence, driven by Guardiola’s tactical expertise and the players’ execution on the pitch. While there are minor areas for improvement, the overall outlook is positive, and they are poised to compete for titles this season. As the season progresses, fans and pundits alike will watch with anticipation to see if City can translate their current form into silverware.



Crew Result:
--------------------------------------------------------------------------------
**Manchester City: A Blueprint of Excellence in Current Form Under Guardiola**

In the ever-dynamic world of Premier League football, Manchester City’s form this season stands as a testament to their consistent excellence and tactical acuity under the stewardship of Pep Guardiola. As the footballing world watches, City’s recent performances have not only reasserted their domestic dominance but have also sent a signal of their intent on the European stage.

Pep Guardiola, the cerebral architect behind Manchester City’s success, recently voiced his satisfaction with the team’s current standing. "The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased," he remarked, encapsulating the harmony and precision that have become synonymous with his tenure at the club.

Historically, Guardiola's sides have been renowned for their possession-based play and tactical flexibility. This season, City continues to exemplify these principles, maintaining an impressive average possession of over 60% in recent matches. This ability to control the game’s tempo has been pivotal in their recent string of victories and is indicative of their strategic dominance on the pitch.

Key to this success are the stellar performances from City’s roster of world-class talent. Kevin De Bruyne, with his unparalleled vision and precision passing, has been instrumental in orchestrating play and dismantling defenses. Erling Haaland, a new addition to the squad, has bolstered City’s attacking threat with his prolific goal-scoring ability, adding a new dimension to their already potent offensive setup. Meanwhile, Phil Foden’s versatility and creativity on the field provide Guardiola with a plethora of tactical options, making him an invaluable asset.

Despite the glowing appraisal, Guardiola remains pragmatic, acknowledging areas for refinement. Defensive lapses, often due to individual errors or brief lapses in concentration, have been points of concern. Addressing these issues is crucial for maintaining their momentum and ensuring that they remain competitive across all fronts.

Guardiola’s comments reflect not only a satisfied manager but also a strategic thinker constantly innovating to maximize the team’s strengths. His ability to adapt formations and player roles is a hallmark of his managerial prowess, ensuring that City not only plays to their strengths but also exploits the weaknesses of their opponents.

Looking ahead, Manchester City faces a challenging fixture schedule, with both domestic and European competitions testing their resolve. However, with Guardiola at the helm and the squad’s depth and quality, they are well-positioned to surmount these challenges. The combination of strategic leadership and player execution suggests that City will continue to be a formidable force in football.

In conclusion, Manchester City’s current form is a masterclass in footballing excellence, driven by Guardiola’s tactical expertise and the players’ execution on the pitch. While there are minor areas for improvement, the overall outlook is positive, and they are poised to compete for titles this season. As the season progresses, fans and pundits alike will watch with anticipation to see if City can translate their current form into silverware.
--------------------------------------------------------------------------------

# Wait for memories to be stored
time.sleep(2)

# List all documents in the collection
try:
    # Query to fetch all documents of this memory type
    query_str = f"SELECT META().id, * FROM `{storage.bucket_name}`.`{storage.scope_name}`.`{storage.collection_name}` WHERE memory_type = $type"
    query_result = storage.cluster.query(query_str, type=storage.type)
    
    print(f"\nAll memory entries in Couchbase:")
    print("-" * 80)
    for i, row in enumerate(query_result, 1):
        doc_id = row.get('id')
        memory_id = row.get(storage.collection_name, {}).get('memory_id', 'unknown')
        content = row.get(storage.collection_name, {}).get('text', '')[:100] + "..."  # Truncate for readability
        
        print(f"Entry {i}:")
        print(f"ID: {doc_id}")
        print(f"Memory ID: {memory_id}")
        print(f"Content: {content}")
        print("-" * 80)
except Exception as e:
    print(f"Failed to list memory entries: {str(e)}")

# Test memory retention
memory_query = "What is Manchester City's current form according to Guardiola?"
memory_results = storage.search(
    query=memory_query,
    limit=5,  # Increased to see more results
    score_threshold=0.0  # Lower threshold to see all results
)

print("\nMemory Search Results:")
print("-" * 80)
for result in memory_results:
    print(f"Context: {result['context']}")
    print(f"Score: {result['score']}")
    print("-" * 80)

# Try a more specific query to find agent interactions
interaction_query = "Manchester City playing style analysis tactical"
interaction_results = storage.search(
    query=interaction_query,
    limit=5,
    score_threshold=0.0
)

print("\nAgent Interaction Memory Results:")
print("-" * 80)
for result in interaction_results:
    print(f"Context: {result['context'][:200]}...")  # Limit output size
    print(f"Score: {result['score']}")
    print("-" * 80)

All memory entries in Couchbase:
--------------------------------------------------------------------------------

Memory Search Results:
--------------------------------------------------------------------------------
Context: Analyzing Manchester City's recent performance based on Pep Guardiola's comments, it is clear that the team is operating at a high level. Guardiola's statement, "The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased," gives insight into both the current form and the overall environment within the squad.

Historically, Manchester City under Guardiola has been known for their possession-based style, high pressing, and tactical flexibility. Comparing their current form to past performances, City appears to be continuing their trend of dominance, especially in terms of controlling matches and creating goal-scoring opportunities. Recent match statistics highlight their effectiveness; for example, in their last few games, City has consistently maintained possession percentages above 60%, showcasing their ability to control the tempo of the game.

Key players such as Kevin De Bruyne, Erling Haaland, and Phil Foden have been pivotal in their recent successes. De Bruyne's vision and passing ability have been instrumental in breaking down defenses, while Haaland's goal-scoring prowess has added a new dimension to City's attacking threat. Foden's versatility and creativity have also been crucial, providing Guardiola with multiple tactical options.

Despite the positive comments from Guardiola, there are still areas where City can improve. Defensive lapses have occasionally been a concern, as seen in matches where they conceded goals due to individual errors or lapses in concentration. Addressing these issues will be critical for maintaining their form throughout the season.

Guardiola's satisfaction with the team's training and performance indicates a harmonious and focused squad, which is essential for sustained success. His comments also reflect the strategic innovations he has implemented, such as adapting formations and player roles to maximize the team's strengths and exploit opponents' weaknesses.

In conclusion, Manchester City's current form is a testament to Guardiola's tactical expertise and the players' execution on the pitch. While there are minor areas for improvement, the overall outlook is positive, and they are well-positioned to compete for titles this season. The combination of Guardiola's strategic leadership and the team's depth and quality suggests that City will continue to be a formidable force in football.
Score: 0.7799903154373169
--------------------------------------------------------------------------------
Context: Pep Guardiola praised Manchester City's current form, saying 'The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased.'
Score: 0.7145987749099731
--------------------------------------------------------------------------------
Context: **Manchester City: A Blueprint of Excellence in Current Form Under Guardiola**

In the ever-dynamic world of Premier League football, Manchester City’s form this season stands as a testament to their consistent excellence and tactical acuity under the stewardship of Pep Guardiola. As the footballing world watches, City’s recent performances have not only reasserted their domestic dominance but have also sent a signal of their intent on the European stage.

Pep Guardiola, the cerebral architect behind Manchester City’s success, recently voiced his satisfaction with the team’s current standing. "The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased," he remarked, encapsulating the harmony and precision that have become synonymous with his tenure at the club.

Historically, Guardiola's sides have been renowned for their possession-based play and tactical flexibility. This season, City continues to exemplify these principles, maintaining an impressive average possession of over 60% in recent matches. This ability to control the game’s tempo has been pivotal in their recent string of victories and is indicative of their strategic dominance on the pitch.

Key to this success are the stellar performances from City’s roster of world-class talent. Kevin De Bruyne, with his unparalleled vision and precision passing, has been instrumental in orchestrating play and dismantling defenses. Erling Haaland, a new addition to the squad, has bolstered City’s attacking threat with his prolific goal-scoring ability, adding a new dimension to their already potent offensive setup. Meanwhile, Phil Foden’s versatility and creativity on the field provide Guardiola with a plethora of tactical options, making him an invaluable asset.

Despite the glowing appraisal, Guardiola remains pragmatic, acknowledging areas for refinement. Defensive lapses, often due to individual errors or brief lapses in concentration, have been points of concern. Addressing these issues is crucial for maintaining their momentum and ensuring that they remain competitive across all fronts.

Guardiola’s comments reflect not only a satisfied manager but also a strategic thinker constantly innovating to maximize the team’s strengths. His ability to adapt formations and player roles is a hallmark of his managerial prowess, ensuring that City not only plays to their strengths but also exploits the weaknesses of their opponents.

Looking ahead, Manchester City faces a challenging fixture schedule, with both domestic and European competitions testing their resolve. However, with Guardiola at the helm and the squad’s depth and quality, they are well-positioned to surmount these challenges. The combination of strategic leadership and player execution suggests that City will continue to be a formidable force in football.

In conclusion, Manchester City’s current form is a masterclass in footballing excellence, driven by Guardiola’s tactical expertise and the players’ execution on the pitch. While there are minor areas for improvement, the overall outlook is positive, and they are poised to compete for titles this season. As the season progresses, fans and pundits alike will watch with anticipation to see if City can translate their current form into silverware.
Score: 0.6959875822067261
--------------------------------------------------------------------------------

Agent Interaction Memory Results:
--------------------------------------------------------------------------------
Context: Analyzing Manchester City's recent performance based on Pep Guardiola's comments, it is clear that the team is operating at a high level. Guardiola's statement, "The team is playing well, we are in a ...
Score: 0.6829519271850586
--------------------------------------------------------------------------------
Context: **Manchester City: A Blueprint of Excellence in Current Form Under Guardiola**

In the ever-dynamic world of Premier League football, Manchester City’s form this season stands as a testament to their ...
Score: 0.6269127130508423
--------------------------------------------------------------------------------
Context: Pep Guardiola praised Manchester City's current form, saying 'The team is playing well, we are in a good moment. The way we are training, the way we are playing - I am really pleased.'...
Score: 0.5217782855033875
--------------------------------------------------------------------------------