Working with Knowledge Bases

from sentence_transformers import SentenceTransformer
import chromadb

class TweetKnowledgeBase:
    def __init__(self, collection_name: str = "twitter_knowledge"):
        # Initialize ChromaDB with persistence
        self.client = chromadb.PersistentClient(path="./chroma_db")
        
        # Use advanced embedding model
        self.embedding_model = SentenceTransformer('all-mpnet-base-v2')
        
        # Create embedding function
        class EmbeddingFunction:
            def __init__(self, model):
                self.model = model
            
            def __call__(self, input: List[str]) -> List[List[float]]:
                embeddings = self.model.encode(input)
                return embeddings.tolist()
        
        embedding_func = EmbeddingFunction(self.embedding_model)
        
        # Create or get collection
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=embedding_func
        )

    def add_tweets(self, tweets: List[Tweet]):
        """Add tweets to the knowledge base."""
        documents = [tweet.text for tweet in tweets]
        ids = [tweet.id for tweet in tweets]
        metadata = [
            {
                "author_id": tweet.author_id,
                "created_at": tweet.created_at,
            }
            for tweet in tweets
        ]
        
        self.collection.add(
            documents=documents,
            ids=ids,
            metadatas=metadata
        )

    def query_knowledge_base(self, query: str, n_results: int = 10) -> List[Dict]:
        """Query the knowledge base for relevant tweets."""
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )
        
        formatted_results = []
        for doc, metadata, distance in zip(
            results['documents'][0], 
            results['metadatas'][0],
            results['distances'][0]
        ):
            formatted_results.append({
                "text": doc,
                "metadata": metadata,
                "relevance_score": 1 - distance  # Convert distance to similarity
            })
        
        return sorted(formatted_results, key=lambda x: x['relevance_score'], reverse=True)

async def update_knowledge_base(twitter_client: TwitterClient, knowledge_base, kol_list: List[Dict]):
    """Update the knowledge base with recent tweets from top KOLs."""
    TOP_KOLS = 5
    TWEETS_PER_KOL = 15
    REQUEST_DELAY = 5
    
    # Select random sample of KOLs
    selected_kols = random.sample(kol_list, min(TOP_KOLS, len(kol_list)))
    
    # Clear existing knowledge base
    knowledge_base.clear_collection()
    
    # Process each selected KOL
    for kol in selected_kols:
        tweets = await twitter_client.get_user_tweets(
            user_id=kol['user_id'],
            max_results=TWEETS_PER_KOL
        )
        
        if tweets:
            knowledge_base.add_tweets(tweets)
        
        # Rate limiting
        await asyncio.sleep(REQUEST_DELAY)

if os.getenv("USE_TWITTER_KNOWLEDGE_BASE", "true").lower() == "true" and knowledge_base is not None:
    tools.append(Tool(
        name="query_twitter_knowledge_base",
        description="""Query the Twitter knowledge base for insights from key opinion leaders.
        Returns relevant tweets that match your query.
        Use this to understand current discussions, trends, and opinions from influential accounts.
        
        Example: query_twitter_knowledge_base("latest developments in AI")
        """,
        func=lambda query: knowledge_base.query_knowledge_base(query)
    ))

from pydantic import BaseModel

class PodcastSegment(BaseModel):
    id: str
    speaker: str
    content: str
    source_file: str
    timestamp: str = None

class PodcastKnowledgeBase:
    def __init__(self, collection_name: str = "podcast_knowledge"):
        # Initialize ChromaDB with persistence
        self.client = chromadb.PersistentClient(path="./chroma_db")
        
        # Use same embedding model as Twitter KB
        self.embedding_model = SentenceTransformer('all-mpnet-base-v2')
        
        # Create collection with embedding function
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=embedding_func
        )

    def process_json_file(self, file_path: str):
        """Process a podcast transcript JSON file."""
        with open(file_path, 'r', encoding='utf-8') as f:
            transcript_data = json.load(f)
        
        segments = []
        for idx, entry in enumerate(transcript_data):
            segment = PodcastSegment(
                id=f"{os.path.basename(file_path)}_{idx}",
                speaker=entry['speaker'],
                content=entry['content'],
                source_file=file_path
            )
            segments.append(segment)
        
        self.add_segments(segments)

    def process_all_json_files(self, directory: str = "youtube_scraper/jsonoutputs"):
        """Process all JSON files, skipping already processed ones."""
        json_files = [f for f in os.listdir(directory) if f.endswith('.json')]
        processed_files = self.get_processed_files()
        
        # Filter out already processed files
        new_files = [f for f in json_files if f not in processed_files]
        
        for json_file in new_files:
            file_path = os.path.join(directory, json_file)
            self.process_json_file(file_path)

if os.getenv("USE_PODCAST_KNOWLEDGE_BASE", "true").lower() == "true" and podcast_knowledge_base is not None:
    tools.append(Tool(
        name="query_podcast_knowledge_base",
        func=lambda query: podcast_knowledge_base.format_query_results(
            podcast_knowledge_base.query_knowledge_base(query)
        ),
        description="""Query the podcast knowledge base for information from podcast transcripts.
        Returns relevant segments from podcast episodes.
        Use this to answer questions about topics discussed in the podcast.
        
        Example: query_podcast_knowledge_base("What did the guest say about DeFi?")
        """
    ))

async def generate_llm_podcast_query(llm: ChatAnthropic = None) -> str:
    """Generates a dynamic, contextually-aware query for the podcast knowledge base."""
    llm = ChatAnthropic(model="claude-3-5-haiku-20241022")
    
    # Format the prompt with random selections
    prompt = PODCAST_QUERY_PROMPT.format(
        topics=random.sample(PODCAST_TOPICS, 3),
        aspects=random.sample(PODCAST_ASPECTS, 2)
    )
    
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    query = response.content.strip()
    
    return query

import chromadb
from sentence_transformers import SentenceTransformer
from typing import List, Dict
from pydantic import BaseModel

class Document(BaseModel):
    id: str
    content: str
    metadata: Dict[str, str]

class MyKnowledgeBase:
    def __init__(self, collection_name: str = "my_knowledge"):
        self.client = chromadb.PersistentClient(path="./chroma_db")
        self.embedding_model = SentenceTransformer('all-mpnet-base-v2')
        
        class EmbeddingFunction:
            def __init__(self, model):
                self.model = model
            def __call__(self, input: List[str]) -> List[List[float]]:
                return self.model.encode(input).tolist()
        
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=EmbeddingFunction(self.embedding_model)
        )
    
    def add_documents(self, documents: List[Document]):
        self.collection.add(
            documents=[doc.content for doc in documents],
            ids=[doc.id for doc in documents],
            metadatas=[doc.metadata for doc in documents]
        )
    
    def query(self, query: str, n_results: int = 5) -> List[Dict]:
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )
        
        return [{
            "content": doc,
            "metadata": meta,
            "score": 1 - dist
        } for doc, meta, dist in zip(
            results['documents'][0],
            results['metadatas'][0],
            results['distances'][0]
        )]

from my_agent.my_knowledge_base import MyKnowledgeBase

def create_agent_tools(llm, knowledge_base, podcast_knowledge_base, agent_kit, config):
    tools = []
    
    if os.getenv("USE_MY_KB", "false").lower() == "true":
        my_kb = MyKnowledgeBase()
        # Load your data
        my_kb.add_documents(load_my_documents())
        
        tools.append(Tool(
            name="query_my_knowledge_base",
            func=lambda q: my_kb.query(q),
            description="Query my custom knowledge base"
        ))
    
    return tools