langchain/templates/rag-codellama-fireworks/rag_codellama_fireworks/chain.py

import os

from git import Repo
from langchain.document_loaders.generic import GenericLoader
from langchain.document_loaders.parsers import LanguageParser
from langchain.embeddings import GPT4AllEmbeddings
from langchain.llms.fireworks import Fireworks
from langchain.prompts import ChatPromptTemplate
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.pydantic_v1 import BaseModel
from langchain_core.runnables import RunnableParallel, RunnablePassthrough

# Check API key
if os.environ.get("FIREWORKS_API_KEY", None) is None:
    raise Exception("Missing `FIREWORKS_API_KEY` environment variable.")

# Load codebase
# Set local path
repo_path = "/Users/rlm/Desktop/tmp_repo"
# Use LangChain as an example
repo = Repo.clone_from("https://github.com/langchain-ai/langchain", to_path=repo_path)
loader = GenericLoader.from_filesystem(
    repo_path + "/libs/langchain/langchain",
    glob="**/*",
    suffixes=[".py"],
    parser=LanguageParser(language=Language.PYTHON, parser_threshold=500),
)
documents = loader.load()

# Split
python_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.PYTHON, chunk_size=2000, chunk_overlap=200
)
texts = python_splitter.split_documents(documents)

# Add to vectorDB
vectorstore = Chroma.from_documents(
    documents=texts,
    collection_name="codebase-rag",
    embedding=GPT4AllEmbeddings(),
)
retriever = vectorstore.as_retriever()

# RAG prompt
template = """Answer the question based only on the following context:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)

# Initialize a Fireworks model
model = Fireworks(model="accounts/fireworks/models/llama-v2-34b-code-instruct")

# RAG chain
chain = (
    RunnableParallel({"context": retriever, "question": RunnablePassthrough()})
    | prompt
    | model
    | StrOutputParser()
)


# Add typing for input
class Question(BaseModel):
    __root__: str


chain = chain.with_types(input_type=Question)
Codebase RAG fireworks (#12597) 11 months ago			`import os`

			`from git import Repo`
			`from langchain.document_loaders.generic import GenericLoader`
			`from langchain.document_loaders.parsers import LanguageParser`
			`from langchain.embeddings import GPT4AllEmbeddings`
			`from langchain.llms.fireworks import Fireworks`
			`from langchain.prompts import ChatPromptTemplate`
			`from langchain.text_splitter import Language, RecursiveCharacterTextSplitter`
			`from langchain.vectorstores import Chroma`
docs[patch], templates[patch]: Import from core (#14575) Update imports to use core for the low-hanging fruit changes. Ran following ```bash git grep -l 'langchain.schema.runnable' {docs,templates,cookbook} \| xargs sed -i '' 's/langchain\.schema\.runnable/langchain_core.runnables/g' git grep -l 'langchain.schema.output_parser' {docs,templates,cookbook} \| xargs sed -i '' 's/langchain\.schema\.output_parser/langchain_core.output_parsers/g' git grep -l 'langchain.schema.messages' {docs,templates,cookbook} \| xargs sed -i '' 's/langchain\.schema\.messages/langchain_core.messages/g' git grep -l 'langchain.schema.chat_histry' {docs,templates,cookbook} \| xargs sed -i '' 's/langchain\.schema\.chat_history/langchain_core.chat_history/g' git grep -l 'langchain.schema.prompt_template' {docs,templates,cookbook} \| xargs sed -i '' 's/langchain\.schema\.prompt_template/langchain_core.prompts/g' git grep -l 'from langchain.pydantic_v1' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.pydantic_v1/from langchain_core.pydantic_v1/g' git grep -l 'from langchain.tools.base' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.tools\.base/from langchain_core.tools/g' git grep -l 'from langchain.chat_models.base' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.chat_models.base/from langchain_core.language_models.chat_models/g' git grep -l 'from langchain.llms.base' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.llms\.base\ /from langchain_core.language_models.llms\ /g' git grep -l 'from langchain.embeddings.base' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.embeddings\.base/from langchain_core.embeddings/g' git grep -l 'from langchain.vectorstores.base' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.vectorstores\.base/from langchain_core.vectorstores/g' git grep -l 'from langchain.agents.tools' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.agents\.tools/from langchain_core.tools/g' git grep -l 'from langchain.schema.output' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.schema\.output\ /from langchain_core.outputs\ /g' git grep -l 'from langchain.schema.embeddings' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.schema\.embeddings/from langchain_core.embeddings/g' git grep -l 'from langchain.schema.document' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.schema\.document/from langchain_core.documents/g' git grep -l 'from langchain.schema.agent' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.schema\.agent/from langchain_core.agents/g' git grep -l 'from langchain.schema.prompt ' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.schema\.prompt\ /from langchain_core.prompt_values /g' git grep -l 'from langchain.schema.language_model' {docs,templates,cookbook} \| xargs sed -i '' 's/from langchain\.schema\.language_model/from langchain_core.language_models/g' ``` 9 months ago			`from langchain_core.output_parsers import StrOutputParser`
			`from langchain_core.pydantic_v1 import BaseModel`
			`from langchain_core.runnables import RunnableParallel, RunnablePassthrough`
Codebase RAG fireworks (#12597) 11 months ago
			`# Check API key`
			`if os.environ.get("FIREWORKS_API_KEY", None) is None:`
			raise Exception("Missing `FIREWORKS_API_KEY` environment variable.")

			`# Load codebase`
			`# Set local path`
			`repo_path = "/Users/rlm/Desktop/tmp_repo"`
			`# Use LangChain as an example`
			`repo = Repo.clone_from("https://github.com/langchain-ai/langchain", to_path=repo_path)`
			`loader = GenericLoader.from_filesystem(`
			`repo_path + "/libs/langchain/langchain",`
			`glob="*/",`
			`suffixes=[".py"],`
			`parser=LanguageParser(language=Language.PYTHON, parser_threshold=500),`
			`)`
			`documents = loader.load()`

			`# Split`
			`python_splitter = RecursiveCharacterTextSplitter.from_language(`
			`language=Language.PYTHON, chunk_size=2000, chunk_overlap=200`
			`)`
			`texts = python_splitter.split_documents(documents)`

			`# Add to vectorDB`
			`vectorstore = Chroma.from_documents(`
			`documents=texts,`
			`collection_name="codebase-rag",`
			`embedding=GPT4AllEmbeddings(),`
			`)`
			`retriever = vectorstore.as_retriever()`

			`# RAG prompt`
			`template = """Answer the question based only on the following context:`
			`{context}`

			`Question: {question}`
			`"""`
			`prompt = ChatPromptTemplate.from_template(template)`

			`# Initialize a Fireworks model`
			`model = Fireworks(model="accounts/fireworks/models/llama-v2-34b-code-instruct")`

			`# RAG chain`
			`chain = (`
			`RunnableParallel({"context": retriever, "question": RunnablePassthrough()})`
			`\| prompt`
			`\| model`
			`\| StrOutputParser()`
			`)`


			`# Add typing for input`
			`class Question(BaseModel):`
			`__root__: str`


			`chain = chain.with_types(input_type=Question)`