Upload 9 files

callbacks.py

@@ -0,0 +1,24 @@
+
+from typing import Dict, Any, List
+from langchain_core.callbacks import BaseCallbackHandler
+import schemas
+import crud
+
+
+class LogResponseCallback(BaseCallbackHandler):
+
+    def __init__(self, user_request: schemas.UserRequest, db):
+        super().__init__()
+        self.user_request = user_request
+        self.db = db
+
+    def on_llm_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
+        """Run when chain ends running."""
+        message = schemas.MessageBase(message=outputs.generations[0][0].text, type='AI')
+        crud.add_message(self.db, message, self.user_request.username)
+
+    def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> Any:
+        for prompt in prompts:
+            print(prompt)

chains.py

@@ -0,0 +1,54 @@
+import os
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_core.runnables import RunnablePassthrough
+import schemas
+from prompts import (
+    raw_prompt,
+    raw_prompt_formatted, 
+    history_prompt_formatted, 
+    standalone_prompt_formatted,
+    rag_prompt_formatted,
+    format_context,
+    tokenizer
+)
+from data_indexing import DataIndexer
+
+
+data_indexer = DataIndexer()
+
+
+llm = HuggingFaceEndpoint(
+    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
+    huggingfacehub_api_token=os.environ['HF_TOKEN'],
+    max_new_tokens=512,
+    stop_sequences=[tokenizer.eos_token],
+    streaming=True,
+)
+
+simple_chain = (raw_prompt | llm).with_types(input_type=schemas.UserQuestion)
+
+formatted_chain = (
+    raw_prompt_formatted
+    | llm
+).with_types(input_type=schemas.UserQuestion)
+
+history_chain = (
+    history_prompt_formatted
+    | llm
+).with_types(input_type=schemas.HistoryInput)
+
+rag_chain = (
+    RunnablePassthrough.assign(new_question=standalone_prompt_formatted | llm)
+    | {
+        'context': lambda x: format_context(data_indexer.search(x['new_question'], hybrid_search=x['hybrid_search'])),
+        'standalone_question': lambda x: x['new_question'],
+        'test': lambda x : print(x)
+    }
+    | rag_prompt_formatted
+    | llm
+).with_types(input_type=schemas.RagInput)
+
+
+
+
+

crud.py

@@ -0,0 +1,27 @@
+from sqlalchemy.orm import Session
+import models, schemas
+
+
+def get_or_create_user(db: Session, username: str):
+    user = db.query(models.User).filter(models.User.username == username).first()
+    if not user:
+        user = models.User(username=username)
+        db.add(user)
+        db.commit()
+        db.refresh(user)
+    return user
+
+def add_message(db: Session, message: schemas.MessageBase, username: str):
+    user = get_or_create_user(db, username)
+    message = models.Message(**message.dict())
+    message.user = user
+    db.add(message)
+    db.commit()
+    db.refresh(message)
+    return message
+
+def get_user_chat_history(db: Session, username: str):
+    user = db.query(models.User).filter(models.User.username == username).first()
+    if not user:
+        return []
+    return user.messages

data_indexing.py

@@ -0,0 +1,154 @@
+import os
+import uuid
+from pathlib import Path
+from pinecone.grpc import PineconeGRPC as Pinecone
+from pinecone import ServerlessSpec
+from langchain_community.vectorstores import Chroma
+from langchain_openai import OpenAIEmbeddings
+
+current_dir = Path(__file__).resolve().parent
+
+
+os.environ['PINECONE_API_KEY'] = "988da8ab-3725-4047-b622-cc42d07ecb6c"
+os.environ['OPENAI_API_KEY'] = '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'
+
+
+class DataIndexer:
+
+    source_file =  os.path.join(current_dir, 'sources.txt')
+
+    def __init__(self, index_name='langchain-repo') -> None:
+        # self.embedding_client = InferenceClient(
+        #     "dunzhang/stella_en_1.5B_v5",
+        # )
+        self.embedding_client = OpenAIEmbeddings()
+
+        self.index_name = index_name
+        self.pinecone_client = Pinecone(api_key=os.environ.get('PINECONE_API_KEY'))
+
+        if index_name not in self.pinecone_client.list_indexes().names():
+            self.pinecone_client.create_index(
+                name=index_name,
+                dimension=1536,
+                metric='cosine',
+                spec=ServerlessSpec(
+                    cloud='aws', 
+                    region='us-east-1'
+                ) 
+            ) 
+
+        self.index = self.pinecone_client.Index(self.index_name)
+        # self.source_index = self.get_source_index()
+        self.source_index = None
+
+    def get_source_index(self):
+        if not os.path.isfile(self.source_file):
+            print('No source file')
+            return None
+        
+        print('create source index')
+        
+        with open(self.source_file, 'r') as file:
+            sources = file.readlines()
+            
+        sources = [s.rstrip('\n') for s in sources]
+        vectorstore = Chroma.from_texts(
+            sources, embedding=self.embedding_client
+        )
+        return vectorstore
+
+    def index_data(self, docs, batch_size=32):
+
+        with open(self.source_file, 'a') as file:
+            for doc in docs:
+                file.writelines(doc.metadata['source'] + '\n')
+
+        for i in range(0, len(docs), batch_size):
+            batch = docs[i: i + batch_size]
+            values = self.embedding_client.embed_documents([
+                doc.page_content for doc in batch
+            ])
+
+            # values = self.embedding_client.feature_extraction([
+            #     doc.page_content for doc in batch
+            # ])
+            vector_ids = [str(uuid.uuid4()) for _ in batch]
+ 
+            metadatas = [{
+                'text': doc.page_content,
+                **doc.metadata
+            } for doc in batch]
+
+            vectors = [{
+                'id': vector_id,
+                'values': value,
+                'metadata': metadata
+            } for vector_id, value, metadata in zip(vector_ids, values, metadatas)]
+
+            try: 
+                upsert_response = self.index.upsert(vectors=vectors)
+                print(upsert_response)
+            except Exception as e:
+                print(e)
+
+    def search(self, text_query, top_k=5, hybrid_search=False):
+
+        print('text query:', text_query)
+
+        filter = None
+        if hybrid_search and self.source_index:
+            source_docs = self.source_index.similarity_search(text_query, 50)
+            print("source_docs", source_docs)
+            filter = {"source": {"$in":[doc.page_content for doc in source_docs]}}
+
+        # vector = self.embedding_client.feature_extraction(text_query)
+        vector = self.embedding_client.embed_query(text_query)
+        result = self.index.query(
+            vector=vector,
+            top_k=top_k,
+            include_metadata=True,
+            filter=filter
+        )
+
+        docs = []
+        for res in result["matches"]:
+            metadata = res["metadata"]
+            if 'text' in metadata:
+                text = metadata.pop('text')
+                docs.append(text)
+        return docs
+    
+
+if __name__ == '__main__':
+
+    from langchain_community.document_loaders import GitLoader
+    from langchain_text_splitters import (
+        Language,
+        RecursiveCharacterTextSplitter,
+    )
+
+    loader = GitLoader(
+        clone_url="https://github.com/langchain-ai/langchain",
+        repo_path="./code_data/langchain_repo/",
+        branch="master",
+    )
+
+    python_splitter = RecursiveCharacterTextSplitter.from_language(
+        language=Language.PYTHON, chunk_size=10000, chunk_overlap=100
+    )
+
+    docs = loader.load()
+    docs = [doc for doc in docs if doc.metadata['file_type'] in ['.py', '.md']]
+    docs = [doc for doc in docs if len(doc.page_content) < 50000]
+    docs = python_splitter.split_documents(docs)
+    for doc in docs:
+        doc.page_content = '# {}\n\n'.format(doc.metadata['source']) + doc.page_content
+
+    indexer = DataIndexer()
+    with open('/app/sources.txt', 'a') as file:
+        for doc in docs:
+            file.writelines(doc.metadata['source'] + '\n')
+    print('DONE')
+    indexer.index_data(docs)
+
+

database.py

@@ -0,0 +1,12 @@
+from sqlalchemy import create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+SQLALCHEMY_DATABASE_URL = "sqlite:///./test.db"
+
+engine = create_engine(
+    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
+)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+Base = declarative_base()

main.py

@@ -0,0 +1,118 @@
+from langchain_core.runnables import Runnable
+from langchain_core.callbacks import BaseCallbackHandler
+from fastapi import FastAPI, Request, Depends
+from sse_starlette.sse import EventSourceResponse
+from sqlalchemy.orm import Session
+from langserve.serialization import WellKnownLCSerializer
+from typing import Any, List
+import crud, models, schemas
+from database import SessionLocal, engine
+from chains import simple_chain, formatted_chain, history_chain, rag_chain
+from prompts import format_chat_history
+from callbacks import LogResponseCallback
+
+
+models.Base.metadata.create_all(bind=engine)
+
+app = FastAPI()
+
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+
+async def generate_stream(input_data: schemas.BaseModel, runnable: Runnable, callbacks: List[BaseCallbackHandler]=[]):
+    for output in runnable.stream(input_data.dict(), config={"callbacks": callbacks}): 
+        data = WellKnownLCSerializer().dumps(output).decode("utf-8")
+        yield {'data': data, "event": "data"} 
+    yield {"event": "end"}
+
+
+@app.get("/")
+def greet_json():
+    return {"Hello": "World!"}
+
+
+@app.post("/simple/stream")
+async def simple_stream(request: Request):
+    data = await request.json()
+    user_question = schemas.UserQuestion(**data['input'])
+    return EventSourceResponse(generate_stream(user_question, simple_chain))
+
+
+@app.post("/formatted/stream")
+async def formatted_stream(request: Request):
+    data = await request.json()
+    user_question = schemas.UserQuestion(**data['input'])
+    return EventSourceResponse(generate_stream(user_question, formatted_chain))
+
+
+@app.post("/history/stream")
+async def history_stream(request: Request, db: Session = Depends(get_db)):  
+    data = await request.json()
+    user_request = schemas.UserRequest(**data['input'])
+    chat_history = crud.get_user_chat_history(db=db, username=user_request.username)
+    message = schemas.MessageBase(message=user_request.question, type='User')
+    crud.add_message(db, message, user_request.username)
+
+    history_input = schemas.HistoryInput(
+        question=user_request.question,
+        chat_history=format_chat_history(chat_history)
+    )
+
+    return EventSourceResponse(generate_stream(
+        history_input, 
+        history_chain,
+        [LogResponseCallback(user_request, db)]
+    ))
+
+
+@app.post("/rag/stream")
+async def rag_stream(request: Request, db: Session = Depends(get_db)):  
+    data = await request.json()
+    user_request = schemas.UserRequest(**data['input'])
+    chat_history = crud.get_user_chat_history(db=db, username=user_request.username)
+    message = schemas.MessageBase(message=user_request.question, type='User')
+    crud.add_message(db, message, user_request.username)
+
+    rag_input = schemas.RagInput(
+        question=user_request.question,
+        chat_history=format_chat_history(chat_history),
+    )
+
+    return EventSourceResponse(generate_stream(
+        rag_input, 
+        rag_chain,
+        [LogResponseCallback(user_request, db)]
+    ))
+
+@app.post("/filtered_rag/stream")
+async def filtered_rag_stream(request: Request, db: Session = Depends(get_db)):  
+    data = await request.json()
+    print(data)
+    user_request = schemas.UserRequest(**data['input'])
+    chat_history = crud.get_user_chat_history(db=db, username=user_request.username)
+    message = schemas.MessageBase(message=user_request.question, type='User')
+    crud.add_message(db, message, user_request.username)
+
+    rag_input = schemas.RagInput(
+        question=user_request.question,
+        chat_history=format_chat_history(chat_history),
+        hybrid_search=True
+    )
+
+    return EventSourceResponse(generate_stream(
+        rag_input, 
+        rag_chain,
+        [LogResponseCallback(user_request, db)]
+    ))
+
+
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("main:app", host="localhost", reload=True,  port=8002)

models.py

@@ -0,0 +1,23 @@
+from sqlalchemy import Column, ForeignKey, Integer, String, DateTime
+from sqlalchemy.orm import relationship
+from datetime import datetime
+
+from database import Base
+
+class User(Base):
+    __tablename__ = "users"
+
+    id = Column(Integer, primary_key=True, index=True)
+    username = Column(String, unique=True, index=True)
+    messages = relationship("Message", back_populates="user")
+
+class Message(Base):
+    __tablename__ = "messages"
+
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(Integer, ForeignKey("users.id"))
+    message = Column(String)
+    type = Column(String)
+    timestamp = Column(DateTime, default=datetime.now)
+
+    user = relationship("User", back_populates="messages")

prompts.py

@@ -0,0 +1,70 @@
+from transformers import AutoTokenizer
+from langchain_core.prompts import PromptTemplate
+from typing import List
+import models
+
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
+
+raw_prompt = "{question}"
+
+history_prompt = """
+Given the following conversation provide a helpful answer to the follow up question.
+
+Chat History:
+{chat_history}
+
+Follow Up question: {question}
+helpful answer:
+"""
+
+standalone_prompt = """
+Given the following conversation and a follow up question, rephrase the 
+follow up question to be a standalone question, in its original language.
+
+Chat History:
+{chat_history}
+
+Follow Up Input: {question}
+
+Standalone question:
+"""
+
+rag_prompt = """
+Answer the question based only on the following context:
+{context}
+
+Question: {standalone_question}
+"""
+
+
+def format_prompt(prompt):
+    chat = [
+        {"role": "system", "content": "You are a helpful AI assistant."},
+        {"role": "user", "content": prompt},
+    ]
+
+    formatted_prompt = tokenizer.apply_chat_template(
+        chat, 
+        tokenize=False, 
+        add_generation_prompt=True
+    )
+
+    return PromptTemplate.from_template(formatted_prompt)
+
+
+def format_chat_history(messages: List[models.Message]):
+    return '\n'.join([
+        '{}: {}'.format(message.type, message.message)
+        for message in messages
+    ])
+
+def format_context(docs: List[str]):
+    return '\n\n'.join(docs)
+
+
+raw_prompt_formatted = format_prompt(raw_prompt)
+raw_prompt = PromptTemplate.from_template(raw_prompt)
+history_prompt_formatted = format_prompt(history_prompt)
+standalone_prompt_formatted = format_prompt(standalone_prompt)
+rag_prompt_formatted = format_prompt(rag_prompt)
+

schemas.py

@@ -0,0 +1,26 @@
+from pydantic.v1 import BaseModel
+from datetime import datetime
+from typing import Optional
+
+class UserQuestion(BaseModel):
+    question: str
+
+class UserRequest(UserQuestion):
+    username: str
+
+class HistoryInput(BaseModel):
+    chat_history: str
+    question: str
+
+class RagInput(HistoryInput):
+    hybrid_search: bool = False
+
+class MessageBase(BaseModel):
+    id: Optional[int] = None
+    user_id: Optional[int] = None
+    message: str
+    type: str
+    timestamp: Optional[datetime] = None
+
+    class Config:
+        orm_mode = True